adjeff 1.2.2__tar.gz → 1.2.3__tar.gz

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

@@ -5,6 +5,19 @@ 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.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.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: adjeff
-Version: 1.2.2
+Version: 1.2.3
 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.3}/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.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "adjeff"
-version = "1.2.2"
+version = "1.2.3"
 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.3}/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.3}/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.3}/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.3}/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.3}/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.3}/src/adjeff/modules/sweep_sampler.py

@@ -78,14 +78,17 @@ 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:
+                    self._dropped.update({str(d): c for d, c in array.coords.items()})
                     arrays[name] = array.drop_vars(list(array.coords), errors="ignore")
         return xr.Dataset(arrays)
 
@@ -96,19 +99,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 +150,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.3}/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.3}/tests/test_atmo_samplers.py

@@ -305,3 +305,32 @@ 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())