Coverage for python/lsst/images/fields/_spline.py: 77%

1# This file is part of lsst-images. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# Use of this source code is governed by a 3-clause BSD-style 

10# license that can be found in the LICENSE file. 

11 

12from __future__ import annotations 

13 

14__all__ = ("SplineField", "SplineFieldSerializationModel") 

15 

16from typing import TYPE_CHECKING, Any, ClassVar, Literal, final 

17 

18import astropy.units 

19import numpy as np 

20import pydantic 

21from scipy.interpolate import Akima1DInterpolator 

22 

23from .._concrete_bounds import BoundsSerializationModel 

24from .._geom import Bounds, Box, Interval 

25from .._image import Image 

26from ..serialization import ( 

27 ArchiveTree, 

28 ArrayReferenceModel, 

29 InlineArray, 

30 InlineArrayModel, 

31 InputArchive, 

32 InvalidParameterError, 

33 NumberType, 

34 OutputArchive, 

35 Unit, 

36) 

37from ._base import BaseField 

38 

39if TYPE_CHECKING: 

40 try: 

41 from lsst.afw.math import BackgroundMI as LegacyBackground 

42 except ImportError: 

43 type LegacyBackground = Any # type: ignore[no-redef] 

44 

45 

46@final 

47class SplineField(BaseField): 

48 """A 2-d Akima spline interpolation of data on a regular grid. 

49 

50 Parameters 

51 ---------- 

52 bounds 

53 The region where this field can be evaluated. 

54 data 

55 The data points to be interpolated. Missing values (indicated by NaN) 

56 are allowed. Will be set to read-only in place. 

57 y 

58 Coordinates for the first dimension of ``data``. Will be set to 

59 read-only in place. 

60 x 

61 Coordinates for the second dimension of ``data``. Will be set to 

62 read-only in place. 

63 unit 

64 Units of the field. 

65 

66 Notes 

67 ----- 

68 This field is much faster to evaluate on a grid via `render` than at 

69 arbitrary points via the function-call operator. 

70 """ 

71 

72 def __init__( 

73 self, 

74 bounds: Bounds, 

75 data: np.ndarray, 

76 *, 

77 y: np.ndarray, 

78 x: np.ndarray, 

79 unit: astropy.units.UnitBase | None = None, 

80 ) -> None: 

81 if isinstance(data, astropy.units.Quantity): 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true

82 if unit is not None: 

83 raise TypeError("If 'data' is a Quantity, 'unit' cannot be provided separately.") 

84 unit = data.unit 

85 data = data.to_value() 

86 if data.ndim != 2: 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true

87 raise ValueError("'data' must be 2-d.") 

88 if y.ndim != 1: 88 ↛ 89line 88 didn't jump to line 89 because the condition on line 88 was never true

89 raise ValueError("'y' must be 1-d.") 

90 if not y.size: 90 ↛ 91line 90 didn't jump to line 91 because the condition on line 90 was never true

91 raise ValueError("No y grid points.") 

92 if not np.all(y[:-1] < y[1:]): 92 ↛ 93line 92 didn't jump to line 93 because the condition on line 92 was never true

93 raise ValueError(f"'y' must be monotonically increasing; got {y}") 

94 if x.ndim != 1: 94 ↛ 95line 94 didn't jump to line 95 because the condition on line 94 was never true

95 raise ValueError("'x' must be 1-d.") 

96 if not x.size: 96 ↛ 97line 96 didn't jump to line 97 because the condition on line 96 was never true

97 raise ValueError("No x grid points.") 

98 if not np.all(x[:-1] < x[1:]): 98 ↛ 99line 98 didn't jump to line 99 because the condition on line 98 was never true

99 raise ValueError(f"'x' must be monotonically increasing; got {x}") 

100 if data.shape != y.shape + x.shape: 100 ↛ 101line 100 didn't jump to line 101 because the condition on line 100 was never true

101 raise ValueError( 

102 f"Shape of 2-d 'data' {data.shape} does not match " 

103 f"expected 1-d 'y' {y.shape} and/or 'x' {x.shape}." 

104 ) 

105 self._bounds = bounds 

106 self._data = data 

107 self._data.flags.writeable = False 

108 self._x = x 

109 self._x.flags.writeable = False 

110 self._y = y 

111 self._y.flags.writeable = False 

112 self._unit = unit 

113 

114 def __eq__(self, other: object) -> bool: 

115 if type(other) is not SplineField: 115 ↛ 116line 115 didn't jump to line 116 because the condition on line 115 was never true

116 return NotImplemented 

117 return ( 

118 self._bounds == other._bounds 

119 and self._unit == other._unit 

120 and np.array_equal(self._data, other._data, equal_nan=True) 

121 and np.array_equal(self._x, other._x, equal_nan=True) 

122 and np.array_equal(self._y, other._y, equal_nan=True) 

123 ) 

124 

125 __hash__ = None # type: ignore[assignment] 

126 

127 @property 

128 def bounds(self) -> Bounds: 

129 return self._bounds 

130 

131 @property 

132 def unit(self) -> astropy.units.UnitBase | None: 

133 return self._unit 

134 

135 @property 

136 def data(self) -> np.ndarray: 

137 """The data points to be interpolated (`numpy.ndarray`). 

138 

139 May have missing values indicated by NaNs. 

140 """ 

141 return self._data 

142 

143 @property 

144 def x(self) -> np.ndarray: 

145 """Coordinates for the second dimension of `data` (`numpy.ndarray`).""" 

146 return self._x 

147 

148 @property 

149 def y(self) -> np.ndarray: 

150 """Coordinates for the first dimension of `data` (`numpy.ndarray`).""" 

151 return self._y 

152 

153 @property 

154 def is_constant(self) -> bool: 

155 # We really do want an exact floating-point comparison here. 

156 return (self._data == self._data[0, 0]).all() 

157 

158 def evaluate( 

159 self, *, x: np.ndarray, y: np.ndarray, quantity: bool = False 

160 ) -> np.ndarray | astropy.units.Quantity: 

161 y, x = np.broadcast_arrays(y, x) 

162 xg = self._x 

163 y_render = np.zeros(xg.shape + y.shape, dtype=np.float64) 

164 mask = np.zeros(xg.size, dtype=bool) 

165 for j in range(xg.size): 

166 if (y_interpolator := self._make_y_interpolator(j)) is not None: 166 ↛ 165line 166 didn't jump to line 165 because the condition on line 166 was always true

167 y_render[j, ...] = y_interpolator(y) 

168 mask[j] = True 

169 if not np.all(mask): 169 ↛ 170line 169 didn't jump to line 170 because the condition on line 169 was never true

170 y_render = y_render[mask, ...] 

171 xg = xg[mask] 

172 result = np.zeros(y.shape, dtype=np.float64) 

173 # There doesn't seem to be a way to avoid looping in Python here; 

174 # maybe someday we'll push this down to a compiled language. 

175 x_interval = self.bounds.bbox.x 

176 for i, xv in np.ndenumerate(x): 

177 if (x_interpolator := self._make_1d_interpolator(xg, y_render[:, *i], x_interval)) is None: 177 ↛ 178line 177 didn't jump to line 178 because the condition on line 177 was never true

178 raise ValueError("No valid data points.") 

179 v = x_interpolator(xv) 

180 result[*i] = v 

181 if quantity: 

182 return astropy.units.Quantity(result, self._unit) 

183 return result 

184 

185 def render(self, bbox: Box | None = None, *, dtype: np.typing.DTypeLike | None = None) -> Image: 

186 if bbox is None: 

187 bbox = self.bounds.bbox 

188 xg = self._x 

189 y_render = np.zeros((xg.size, bbox.y.size), dtype=dtype) 

190 mask = np.zeros(xg.size, dtype=bool) 

191 for j in range(xg.size): # we have to loop, but only over bins, not evaluation points. 

192 if (y_interpolator := self._make_y_interpolator(j)) is not None: 192 ↛ 191line 192 didn't jump to line 191 because the condition on line 192 was always true

193 y_render[j, :] = y_interpolator(bbox.y.arange) 

194 mask[j] = True 

195 if not np.all(mask): 195 ↛ 196line 195 didn't jump to line 196 because the condition on line 195 was never true

196 y_render = y_render[mask, :] 

197 xg = xg[mask] 

198 x_interval = self.bounds.bbox.x 

199 if (x_interpolator := self._make_1d_interpolator(xg, y_render, x_interval)) is None: 199 ↛ 200line 199 didn't jump to line 200 because the condition on line 199 was never true

200 raise ValueError("No valid data points.") 

201 rendered_array = x_interpolator(bbox.x.arange) 

202 return Image(rendered_array.transpose().copy(), bbox=bbox, unit=self._unit, dtype=dtype) 

203 

204 def multiply_constant( 

205 self, factor: float | astropy.units.Quantity | astropy.units.UnitBase 

206 ) -> SplineField: 

207 factor, unit = self._handle_factor_units(factor) 

208 return SplineField(self._bounds, self._data * factor, y=self._y, x=self._x, unit=unit) 

209 

210 def serialize(self, archive: OutputArchive[Any]) -> SplineFieldSerializationModel: 

211 """Serialize the spline field to an output archive. 

212 

213 Parameters 

214 ---------- 

215 archive 

216 Archive to write to. 

217 """ 

218 if self._data.size > 64: 218 ↛ 219line 218 didn't jump to line 219 because the condition on line 218 was never true

219 data = archive.add_array(self._data, name="data") 

220 else: 

221 data = InlineArrayModel( 

222 data=self._data.tolist(), 

223 datatype=NumberType.from_numpy(self._data.dtype), 

224 ) 

225 return SplineFieldSerializationModel( 

226 bounds=self.bounds.serialize(), 

227 data=data, 

228 y=self._y, 

229 x=self._x, 

230 unit=self._unit, 

231 ) 

232 

233 @staticmethod 

234 def _get_archive_tree_type( 

235 pointer_type: type[Any], 

236 ) -> type[SplineFieldSerializationModel]: 

237 """Return the serialization model type for this object for an archive 

238 type that uses the given pointer type. 

239 """ 

240 return SplineFieldSerializationModel 

241 

242 @staticmethod 

243 def from_legacy_background( 

244 legacy_background: LegacyBackground, 

245 bounds: Bounds | None = None, 

246 unit: astropy.units.UnitBase | None = None, 

247 ) -> SplineField: 

248 """Convert from a legacy `lsst.afw.math.BackgroundMI` instance. 

249 

250 Parameters 

251 ---------- 

252 legacy_background 

253 Legacy background object to convert. 

254 bounds 

255 The bounds of the returned field, if they should be different from 

256 the bounding box of ``legacy_background``. 

257 unit 

258 The units of the returned field (`lsst.afw.math.Background` 

259 objects do not know their units). 

260 

261 Notes 

262 ----- 

263 `SplineField.render` and the `lsst.afw` background interpolator both 

264 use Akima splines, but with slightly different boundary conditions. 

265 They should produce equivalent to single-precision round-off error 

266 when evaluated within the region enclosed by bin centers (i.e. where 

267 no extrapolation is necessary) and when there are five or more 

268 points to be interpolated in each row and column. 

269 """ 

270 from lsst.afw.math import ApproximateControl, Interpolate 

271 

272 bg_control = legacy_background.getBackgroundControl() 

273 approx_control = bg_control.getApproximateControl() 

274 stats_image = legacy_background.getStatsImage() 

275 # In the afw background system, "approximate" is the opposite of 

276 # "interpolate", but it also implied Chebyshev since that's the only 

277 # approximation algorithm we every implemented. All of the 

278 # interpolation options are similarly splines, and non-Akima splines 

279 # are *mostly* only used when there aren't enough control points for 

280 # Akima splines. Since SciPy automatically falls back to non-Akima 

281 # splines in those cases (or maybe they're formally a limit of Akima 

282 # splines, I don't know), we just always assume what we get can be 

283 # Akima-spline interpolated by SciPy to good enough approximation with 

284 # what afw would do. 

285 if approx_control.getStyle() != ApproximateControl.UNKNOWN: 285 ↛ 286line 285 didn't jump to line 286 because the condition on line 285 was never true

286 raise TypeError("Legacy background uses Chebyshev approximation, not splines.") 

287 if bg_control.getInterpStyle() == Interpolate.UNKNOWN: 287 ↛ 288line 287 didn't jump to line 288 because the condition on line 287 was never true

288 raise TypeError("Legacy background does not use spline interpolation.") 

289 x = legacy_background.getBinCentersX() 

290 y = legacy_background.getBinCentersY() 

291 return SplineField( 

292 Box.from_legacy(legacy_background.getImageBBox()) if bounds is None else bounds, 

293 stats_image.image.array, 

294 x=x, 

295 y=y, 

296 unit=unit, 

297 ) 

298 

299 def _make_1d_interpolator( 

300 self, loc: np.ndarray, val: np.ndarray, fallback_interval: Interval 

301 ) -> Akima1DInterpolator | None: 

302 match len(loc): 

303 case 0: 303 ↛ 304line 303 didn't jump to line 304 because the pattern on line 303 never matched

304 return None 

305 case 1: 305 ↛ 309line 305 didn't jump to line 309 because the pattern on line 305 never matched

306 # SciPy can handle only two points by downgrading to linear 

307 # interpolation, but it raises if given only one. Mock up 

308 # two for the nearest-neighbor fallback. 

309 return Akima1DInterpolator( 

310 np.array([fallback_interval.min - 0.5, fallback_interval.max + 0.5]), 

311 np.array([val[0], val[0]]), 

312 ) 

313 case _: 

314 return Akima1DInterpolator(loc, val, extrapolate=True) 

315 

316 def _make_y_interpolator(self, j: int) -> Akima1DInterpolator | None: 

317 y = self._y 

318 z = self._data[:, j] 

319 mask = np.isfinite(z) 

320 if not np.all(mask): 

321 y = y[mask] 

322 z = z[mask] 

323 del mask 

324 return self._make_1d_interpolator(y, z, self.bounds.bbox.y) 

325 

326 

327class SplineFieldSerializationModel(ArchiveTree): 

328 """Serialization model for `SplineField`.""" 

329 

330 SCHEMA_NAME: ClassVar[str] = "spline_field" 

331 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

332 MIN_READ_VERSION: ClassVar[int] = 1 

333 PUBLIC_TYPE: ClassVar[type] = SplineField 

334 

335 bounds: BoundsSerializationModel = pydantic.Field( 

336 description=("The region where this field can be evaluated.") 

337 ) 

338 

339 data: ArrayReferenceModel | InlineArrayModel = pydantic.Field( 

340 description="2-d data to interpolate. NaNs indicate missing values." 

341 ) 

342 

343 y: InlineArray = pydantic.Field(description="Row positions of the data points.") 

344 

345 x: InlineArray = pydantic.Field(description="Column positions of the data points.") 

346 

347 unit: Unit | None = pydantic.Field(default=None, description="Units of the field.") 

348 

349 field_type: Literal["SPLINE"] = "SPLINE" 

350 

351 def deserialize(self, archive: InputArchive, **kwargs: Any) -> SplineField: 

352 """Deserialize the spline field from an input archive. 

353 

354 Parameters 

355 ---------- 

356 archive 

357 Archive to read from. 

358 **kwargs 

359 Unsupported keyword arguments are accepted only to provide 

360 better error messages (raising 

361 `.serialization.InvalidParameterError`). 

362 """ 

363 if kwargs: 363 ↛ 364line 363 didn't jump to line 364 because the condition on line 363 was never true

364 raise InvalidParameterError(f"Unrecognized parameters for SplineField: {set(kwargs.keys())}.") 

365 data = ( 

366 np.array(self.data.data, dtype=self.data.datatype.to_numpy()) 

367 if isinstance(self.data, InlineArrayModel) 

368 else archive.get_array(self.data) 

369 ) 

370 return SplineField( 

371 self.bounds.deserialize(), 

372 data, 

373 y=self.y, 

374 x=self.x, 

375 unit=self.unit, 

376 )