adjeff 1.2.2__tar.gz → 1.2.5__tar.gz

{adjeff-1.2.2 → adjeff-1.2.5}/CHANGELOG.md

@@ -5,6 +5,40 @@ All notable changes to this project are documented here.
 The format follows [Keep a Changelog](https://keepachangelog.com/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/).
 
+## [1.2.5]
+
+### Fixed
+
+- **A stray scalar coordinate no longer overwrites a swept axis.** A
+  sweep's dropped coordinates were collected from every coordinate an
+  input array carried, not just the ones indexing its own dimensions.
+  A field read back from a scene that already holds per-wavelength
+  results, e.g. for `fit_psf`'s own forward pass, inherits a scalar
+  `wl` tag from that scene; merged in, it silently replaced the real,
+  one-row-per-band `wl` axis, and the next lookup by wavelength raised
+  `ValueError: dimension 'wl' already exists as a scalar variable`.
+- **A loader's atmosphere and geometry now survive a cache hit.**
+  `ProductLoader` declared a dynamic `output_vars` property to aggregate
+  what its mixins contribute, but `SceneModule`'s caching reads
+  `_output_vars`, a different name it never set; a cache write therefore
+  persisted nothing, and the next read restored an empty scene. The first
+  run of a script looked fine, since `_compute` writes straight into the
+  scene; only a second run, or any read of a pre-existing cache, raised
+  `MissingVariableError` on `aot`, `h`, `vza` and the rest.
+
+## [1.2.3]
+
+### Fixed
+
+- **A radiative sweep over maps keeps the scene's coordinates.** xsweep
+  labels by position the dims it does not sweep, so a scene whose `aot`
+  and `h` are maps got its terms aligned on those labels: every pixel
+  came back NaN but the handful whose index matched one. `load_maja(...,
+  as_map=True, compute_radiatives=True)` could not produce anything.
+- **A sweep restores its own coordinates.** A sampler used to do it
+  itself, in its own `_compute`, which is what the fix above forgot in
+  one of five places. There is now one place.
+
 ## [1.2.2]
 
 ### Fixed

{adjeff-1.2.2 → adjeff-1.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: adjeff
-Version: 1.2.2
+Version: 1.2.5
 Summary: Library that provides utilities to simulate adjacency effects.
 Project-URL: Homepage, https://github.com/walcark/Adjeff
 Project-URL: Issues, https://github.com/walcark/Adjeff/issues

{adjeff-1.2.2 → adjeff-1.2.5}/docs/guide.md

@@ -5,14 +5,15 @@
 3. [Atmosphere and geometry](#3-atmosphere-and-geometry)
 4. [Scene modules and pipelines](#4-scene-modules-and-pipelines)
 5. [Radiative samplers](#5-radiative-samplers)
-6. [PSF models](#6-psf-models)
-7. [Correction](#7-correction)
-8. [Fitting a PSF](#8-fitting-a-psf)
-9. [High-level API](#9-high-level-api)
-10. [Analysis and the `.adjeff` accessor](#10-analysis-and-the-adjeff-accessor)
-11. [Cache](#11-cache)
-12. [Logging](#12-logging)
-13. [Installation details](#13-installation-details)
+6. [Writing a sampler](#6-writing-a-sampler)
+7. [PSF models](#7-psf-models)
+8. [Correction](#8-correction)
+9. [Fitting a PSF](#9-fitting-a-psf)
+10. [High-level API](#10-high-level-api)
+11. [Analysis and the `.adjeff` accessor](#11-analysis-and-the-adjeff-accessor)
+12. [Cache](#12-cache)
+13. [Logging](#13-logging)
+14. [Installation details](#14-installation-details)
 
 ---
 
@@ -180,7 +181,51 @@ data, or Smart-G returns NaN, which the modules turn into an error.
 
 > Worked through in [notebook 03](../notebooks/03-compute-radiative-quantities.ipynb).
 
-## 6. PSF models
+## 6. Writing a sampler
+
+A sampler is a combination of a physics function and a sentence saying how to call it.
+Four declarations are required, and then the base class does the rest: cache, batching,
+deduplication, provenance, and putting the coordinates back.
+
+```python
+class SphAlbSampler(AtmoSampler):
+    _output_vars = ["sph_alb"]
+    contract = "batch(aot, rh, h, href) vec(wl) -> sph_alb(wl)"
+    point_fn = staticmethod(sph_alb)
+    default_n_ph = 20000000
+```
+
+`point_fn` is the physics: it takes one state and returns one result.
+The contract says what a state is.
+
+### The contract
+
+| Word | Meaning |
+|---|---|
+| `vec(...)` | these parameters travel together in one call, the physics gets arrays |
+| `batch(...)` | these parameters are grouped into calls of `batch_size` states |
+| `loop(...)` | one call per value, nothing else is possible |
+| `-> name(dims)` | the variable written, and the dims it comes back with |
+
+### Writing the output
+
+A sampler whose contract carries `wl` is swept once for every band. One
+that reads the scene, like `RhoToaSymSampler`, sweeps per band instead
+and passes the scene to the physics:
+
+```python
+def _compute(self, scene):
+    for band in scene.bands:
+        scene[band][self._slot("rho_toa")] = self._sweep(
+            rho_s=scene[band], band=band
+        )
+    return scene
+```
+
+What is bound this way stays out of the sweep space, so the physics can
+read a map the sweep knows nothing about.
+
+## 7. PSF models
 
 Analytical PSFs are `torch.nn.Module`s with bounded parameters, sampled
 on an odd, square `PSFGrid`.
@@ -213,7 +258,7 @@ psf_params(tree, S2Band.B02)      # {"sigma": DataArray}
 tree.to_zarr("psf.zarr", mode="w")
 ```
 
-## 7. Correction
+## 8. Correction
 
 For a uniform surface, `Unif2Toa` applies the 5S formula and `Toa2Unif`
 inverts it, giving `rho_unif`, the reflectance a uniform surface would
@@ -234,7 +279,7 @@ from adjeff.modules.models import Unif2Surface
 scene = Unif2Surface(kernels=tree)(scene)
 ```
 
-## 8. Fitting a PSF
+## 9. Fitting a PSF
 
 `fit` optimises the PSF of a model over reference scenes, once per
 band and atmospheric state, and returns the frozen tree.
@@ -275,7 +320,7 @@ encircled-energy radii over a list of PSFs, e.g. a parameter grid.
 
 > Worked through in [notebook 05](../notebooks/05-learn-psf.ipynb).
 
-## 9. High-level API
+## 10. High-level API
 
 `adjeff` exports one-call functions for the common paths:
 
@@ -294,7 +339,7 @@ encircled-energy radii over a list of PSFs, e.g. a parameter grid.
 
 > Worked through in [notebook 04](../notebooks/04-simulate-rho-toa.ipynb).
 
-## 10. Analysis and the `.adjeff` accessor
+## 11. Analysis and the `.adjeff` accessor
 
 `adjeff.analysis` quantifies a PSF or a 2-D field:
 
@@ -322,7 +367,7 @@ da.adjeff.untidy()                  # and back, before xr.merge
 
 > Worked through in [notebook 06](../notebooks/06-evaluate-a-psf.ipynb).
 
-## 11. Cache
+## 12. Cache
 
 Every module takes a `CacheStore`. Results are stored in Zarr, keyed by
 the module, its configuration and the provenance of its inputs; an
@@ -340,7 +385,7 @@ scene = pipeline(scene)   # read back, no GPU call
 
 Without a cache, every output stays in RAM.
 
-## 12. Logging
+## 13. Logging
 
 Adjeff prints nothing until asked:
 
@@ -362,7 +407,7 @@ are capped at `warning`, and Python warnings join the stream. Lines are
 Everything is an ordinary `logging` record under `adjeff`, so an
 application's own handlers receive it.
 
-## 13. Installation details
+## 14. Installation details
 
 | Environment | GPU | Contents |
 |---|---|---|

{adjeff-1.2.2 → adjeff-1.2.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "adjeff"
-version = "1.2.2"
+version = "1.2.5"
 authors = [{ name = "walcark", email = "kevin.walcarius@gmail.com" }]
 readme = "README.md"
 description = "Library that provides utilities to simulate adjacency effects."

{adjeff-1.2.2 → adjeff-1.2.5}/src/adjeff/modules/loaders/maja_loader.py

@@ -195,6 +195,7 @@ class MajaLoader(
 
         with rasterio.open(glob_mnt[0]) as src:
             arr = src.read(1).astype(float)
+            arr = np.where(arr < 0.0, 0.0, arr)
 
         if self.as_map:
             return xr.DataArray(

{adjeff-1.2.2 → adjeff-1.2.5}/src/adjeff/modules/loaders/product_loader.py

@@ -116,13 +116,19 @@ class ProductLoader(SceneSource, ABC):
     _BASE_OUTPUT_VARS: ClassVar[list[str]] = ["rho_s"]
 
     @property
-    def output_vars(self) -> list[str]:
-        """Aggregate output slots declared by active mixins."""
+    def _output_vars(self) -> list[str]:  # type: ignore[override]
+        """Aggregate output roles declared by active mixins.
+
+        Roles, not slots: ``SceneModule`` turns these into slots itself,
+        for caching, provenance and the public ``output_vars`` it reads
+        them from.  A plain ``ClassVar`` cannot do this, since which
+        mixins are active is only known once this class is built.
+        """
         roles = list(self._BASE_OUTPUT_VARS)
         for cls in type(self).__mro__:
             contrib: list[str] = cls.__dict__.get("_output_vars_contribution", [])
             roles.extend(contrib)
-        return [self._slot(role) for role in roles]
+        return roles
 
     def __init__(
         self,

{adjeff-1.2.2 → adjeff-1.2.5}/src/adjeff/modules/samplers/_brdf_sampler.py

@@ -126,11 +126,7 @@ class BrdfSampler(AtmoSampler):
         name = self._output_vars[0]
         for band in self.spectral_config.bands:
             ds = scene[band]
-            # Coordinates first: the algebra below aligns by label, and
-            # xsweep numbers a dim it did not sweep from zero, which no
-            # longer matches the scene's own labels.
-            arr = self._restore_coords(raw.sel(wl=band.wl_nm), ds)
-            ds[self._slot(name)] = self._normalise(arr, ds)
+            ds[self._slot(name)] = self._normalise(raw.sel(wl=band.wl_nm), ds)
         return scene
 
     def _t_sun(self, ds: xr.Dataset) -> xr.DataArray:

{adjeff-1.2.2 → adjeff-1.2.5}/src/adjeff/modules/samplers/rho_toa.py

@@ -140,6 +140,6 @@ class RhoToaSampler(SweepSampler):
         # the two travel together rather than through the sweep space.
         for band in scene.bands:
             arr = self._sweep(rho_s=scene[band], band=band)
-            scene[band]["rho_toa"] = self._restore_coords(arr, scene[band])
+            scene[band]["rho_toa"] = arr
 
         return scene

{adjeff-1.2.2 → adjeff-1.2.5}/src/adjeff/modules/samplers/rho_toa_sym.py

@@ -125,6 +125,6 @@ class RhoToaSymSampler(SweepSampler):
         # the two travel together rather than through the sweep space.
         for band in scene.bands:
             arr = self._sweep(rho_s=scene[band], band=band)
-            scene[band][self._slot("rho_toa")] = self._restore_coords(arr, scene[band])
+            scene[band][self._slot("rho_toa")] = arr
 
         return scene

{adjeff-1.2.2 → adjeff-1.2.5}/src/adjeff/modules/scene_module.py

@@ -133,6 +133,9 @@ class SceneModule:
 
             outcome["cached"] = False
             scene = self._compute(scene)
+            for band in scene.bands:
+                ds = scene[band]
+                print(ds)
             self._reject_non_finite(scene)
             self._stamp_provenance(scene, key)
             self._cache.save_vars(key, self._role_view(scene), self._output_vars)
@@ -162,6 +165,7 @@ class SceneModule:
                 values = ds[slot]
                 if values.dtype.kind not in "fc":
                     continue
+                print(values)
                 # Reduced rather than materialised: a lazy output is read
                 # back one chunk at a time instead of all at once.
                 bad = int((~np.isfinite(values)).sum())

{adjeff-1.2.2 → adjeff-1.2.5}/src/adjeff/modules/sweep_sampler.py

@@ -78,14 +78,26 @@ class SweepSampler(SceneModule):
     def _space(self) -> xr.Dataset:
         """Collect the swept parameters from the configs into one Dataset.
 
-        Coordinates are dropped: two arrays sharing a spatial dim would
-        otherwise claim conflicting labels for it.
+        Coordinates are dropped, two arrays sharing a spatial dim would
+        otherwise claim conflicting labels for it, and kept aside for
+        :meth:`_sweep` to put back.
         """
         wanted = set(self._contract.inputs)
         arrays: dict[str, xr.DataArray] = {}
+        self._dropped: dict[str, xr.DataArray] = {}
         for config in self._get_configs():
             for name, array in config._arrays.items():
                 if name in wanted and name not in arrays:
+                    # Only an array's own dimension coordinates: a scalar
+                    # tag of the same name, picked up from an earlier
+                    # pipeline stage sharing this Dataset, is not one.
+                    self._dropped.update(
+                        {
+                            str(d): array.coords[d]
+                            for d in array.dims
+                            if d in array.coords
+                        }
+                    )
                     arrays[name] = array.drop_vars(list(array.coords), errors="ignore")
         return xr.Dataset(arrays)
 
@@ -96,19 +108,29 @@ class SweepSampler(SceneModule):
 
         return coerce(type(self).contract)
 
-    @staticmethod
-    def _restore_coords(arr: xr.DataArray, source: xr.Dataset) -> xr.DataArray:
-        """Give *arr* the coordinates *source* holds for its dims.
+    def _restore_coords(
+        self, arr: xr.DataArray, source: xr.Dataset | None = None
+    ) -> xr.DataArray:
+        """Give *arr* back the coordinates the sweep space was stripped of.
 
         xsweep labels the dims it does not sweep (``y``, ``x``) by
         position; assigned as they are, they would misalign with the
-        scene and leave NaN.
+        scene and leave NaN.  *source* names the scene the physics was
+        bound to, whose own dims are not in the space.
         """
         shared = {
-            str(dim): source.coords[dim]
-            for dim in arr.dims
-            if dim in source.coords and dim in source.dims
+            name: coord
+            for name, coord in getattr(self, "_dropped", {}).items()
+            if name in arr.dims
         }
+        if source is not None:
+            shared.update(
+                {
+                    str(dim): source.coords[dim]
+                    for dim in arr.dims
+                    if dim in source.coords and dim in source.dims
+                }
+            )
         return arr.assign_coords(shared) if shared else arr
 
     def _sweep(self, **bound: Any) -> xr.DataArray:
@@ -137,7 +159,7 @@ class SweepSampler(SceneModule):
         self._log.info("sweep.plan", **plan)
         result = sweeper(space, policy=self._policy(space, bound), **statics)
         out: xr.DataArray = result[self._contract.outputs[0]]
-        return out
+        return self._restore_coords(out, bound.get("rho_s"))
 
     _SPATIAL: ClassVar[frozenset[str]] = frozenset({"x", "y", "x_psf", "y_psf"})
 

{adjeff-1.2.2 → adjeff-1.2.5}/src/adjeff/reference/wu2024.py

@@ -125,7 +125,6 @@ class WuPsfSampler(SweepSampler):
         for band in scene.bands:
             ds = scene[band]
             arr = self._sweep(rho_s=ds, band=band)
-            arr = self._restore_coords(arr, ds)
             # The store allocates from the contract and labels a dimension
             # it never swept by position, so the grid is put back here.
             grid = PSFGrid(res=ds["rho_s"].adjeff.res, n=ds["rho_s"].adjeff.n)

{adjeff-1.2.2 → adjeff-1.2.5}/tests/test_atmo_samplers.py

@@ -305,3 +305,63 @@ class TestSweepPolicy:
             for ds in scenes
         ]
         assert stores[0] != stores[1]
+
+    @pytest.mark.integration
+    def test_a_mapped_sweep_keeps_the_scene_coordinates(self):
+        """Every pixel gets a value, not the few whose index matches a label."""
+        import numpy as np
+
+        from adjeff.api import run_radiatives_from_scene
+        from adjeff.core import gaussian_image_dict
+
+        band, size = S2Band.B02, 5
+        scene = gaussian_image_dict(sigma=1.0, res_km=0.5, bands=[band], n=size)
+        ds = scene[band]
+        rng = np.random.default_rng(0)
+        ds["aot"] = (("y", "x"), rng.uniform(0.05, 0.4, (size, size)))
+        ds["h"] = (("y", "x"), np.zeros((size, size)))
+        for name, value in (
+            ("rh", 50.0),
+            ("href", 2.0),
+            ("sza", 17.0),
+            ("saa", 117.0),
+            ("vza", 3.2),
+            ("vaa", 207.0),
+        ):
+            ds[name] = xr.DataArray([value], dims=[name])
+        scene[band] = ds
+
+        pytest.importorskip("smartg")
+        out = run_radiatives_from_scene(scene, n_bins=2)
+        assert bool(np.isfinite(out[band]["tdir_down"]).all())
+
+    def test_a_stray_scalar_coordinate_does_not_overwrite_a_real_axis(self):
+        """An array carrying a stray same-named scalar coord keeps the real one."""
+        from adjeff.modules.samplers import TdirDownSampler
+
+        spectral = atmo.SpectralConfig.from_bands([S2Band.B8A])
+        geo = atmo.GeoConfig(
+            sza=xr.DataArray([17.0], dims=["sza"]),
+            vza=xr.DataArray([3.0], dims=["vza"]),
+            saa=xr.DataArray([117.0], dims=["saa"]),
+            vaa=xr.DataArray([207.0], dims=["vaa"]),
+        )
+        # aot carries a stray scalar 'wl' coordinate, as if inherited from
+        # a Dataset that already held per-wl radiative terms.
+        aot = xr.DataArray([0.2], dims=["aot"]).assign_coords(wl=865.0)
+        sampler = TdirDownSampler(
+            atmo_config=atmo.AtmoConfig(
+                aot=aot,
+                rh=xr.DataArray([50.0], dims=["rh"]),
+                h=xr.DataArray([0.1], dims=["h"]),
+                href=xr.DataArray([2.0], dims=["href"]),
+                species={"sulphate": 1.0},
+            ),
+            geo_config=geo,
+            spectral_config=spectral,
+            remove_rayleigh=False,
+        )
+        sampler._space()
+        wl = sampler._dropped["wl"]
+        assert wl.dims == ("wl",)
+        assert float(wl.values[0]) == 865.0

{adjeff-1.2.2 → adjeff-1.2.5}/tests/test_loaders.py

@@ -79,3 +79,51 @@ class TestMissingInputs:
 
         with pytest.raises(ImageIOError, match="Sun angles"):
             loader._sza_saa()
+
+
+class TestOutputVarsCaching:
+    """A loader's mixin-contributed variables must round-trip through cache."""
+
+    def test_atmosphere_and_geometry_survive_a_cache_hit(self, tmp_path):
+        """The second call, reading the cache, still carries aot, h, vza..."""
+        from adjeff.core import S2Band
+        from adjeff.modules.loaders.product_loader import (
+            AtmosphereMixin,
+            ElevationMixin,
+            GeometryMixin,
+            ProductLoader,
+        )
+        from adjeff.utils import CacheStore
+
+        class FullLoader(GeometryMixin, AtmosphereMixin, ElevationMixin, ProductLoader):
+            def __init__(self, product_path, bands, res, cache=None):
+                AtmosphereMixin.__init__(self)
+                ProductLoader.__init__(self, product_path, bands, res, cache=cache)
+
+            def reflectance(self, band):
+                return _make_ref()
+
+            def _vza_vaa(self, band):
+                return xr.DataArray([3.0]), xr.DataArray([200.0])
+
+            def _sza_saa(self):
+                return xr.DataArray([17.0]), xr.DataArray([117.0])
+
+            def _aot(self, ref):
+                return xr.DataArray([0.2])
+
+            def _rh(self):
+                return xr.DataArray([50.0])
+
+            def _h(self, ref):
+                return xr.DataArray([0.1])
+
+        cache = CacheStore(tmp_path)
+        product_path = tmp_path / "product"
+        product_path.mkdir()
+        bands = [S2Band.B02]
+
+        for _ in range(2):  # first run computes, second reads the cache
+            scene = FullLoader(product_path, bands, res=0.12, cache=cache)()
+            for role in ("rho_s", "aot", "rh", "href", "h", "vza", "vaa", "sza", "saa"):
+                assert role in scene[S2Band.B02], role