adjeff 1.2.1__tar.gz → 1.2.2__tar.gz

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

@@ -5,6 +5,16 @@ 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.2]
+
+### Fixed
+
+- **The sweep store is named after the scene bound to it.** The scene is
+  passed to the physics outside xsweep's fingerprint, so two training
+  scenes sweeping the same space shared a store and the second read back
+  the first one's results. Introduced in 1.1.0, which is where the store
+  comes from.
+
 ## [1.2.1]
 
 ### Fixed
@@ -34,9 +44,8 @@ PSF fit being the case that motivated it.
 - **A cached sweep is written through to zarr and read back lazily.** It
   applies to the samplers whose output carries the image dims, and only
   when the module has a cache directory; a sweep returning a spectrum per
-  state is still assembled in memory. Two consequences: the three training
-  scenes of a fit share one Smart-G run instead of three, and the result a
-  module hands on is dask-backed rather than numpy.
+  state is still assembled in memory. The result a module hands on is then
+  dask-backed rather than numpy.
 - **`batch_size` controls the batching.** It did not before: xsweep took
   the whole space in one call whatever its value. Measured on a 3999 pixel
   grid, a batch of 24 states costs five times the time of three batches of

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

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

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

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

@@ -13,6 +13,7 @@ import functools
 import hashlib
 from typing import TYPE_CHECKING, Any, Callable, ClassVar
 
+import joblib  # type: ignore[import-untyped]
 import numpy as np
 import xarray as xr
 from xsweep import Sweeper, SweepPolicy
@@ -134,13 +135,13 @@ class SweepSampler(SceneModule):
         if spectral is not None:
             plan["bands"] = len(spectral.bands)
         self._log.info("sweep.plan", **plan)
-        result = sweeper(space, policy=self._policy(space), **statics)
+        result = sweeper(space, policy=self._policy(space, bound), **statics)
         out: xr.DataArray = result[self._contract.outputs[0]]
         return out
 
     _SPATIAL: ClassVar[frozenset[str]] = frozenset({"x", "y", "x_psf", "y_psf"})
 
-    def _policy(self, space: xr.Dataset) -> SweepPolicy:
+    def _policy(self, space: xr.Dataset, bound: dict[str, Any]) -> SweepPolicy:
         """Return the sweep policy.
 
         A sweep whose output carries the image dims is written to zarr as
@@ -154,25 +155,42 @@ class SweepSampler(SceneModule):
         return SweepPolicy(
             batch_size=self.batch_size,
             dedup=self.dedup,
-            store=str(root / "sweeps" / self._store_name(space)),
+            store=str(root / "sweeps" / self._store_name(space, bound)),
             chunks=self._chunks(space),
             load=False,
         )
 
-    def _store_name(self, space: xr.Dataset) -> str:
-        """Return a store directory unique to the sampler, its grid and statics.
+    def _store_name(self, space: xr.Dataset, bound: dict[str, Any]) -> str:
+        """Return a store directory unique to everything the results depend on.
 
-        A store cannot be resumed once anything it was written with has
-        changed, the aerosol species as much as an axis, so both are part
-        of its name.
+        The axes and the statics, which xsweep fingerprints itself, and the
+        scene data bound to the physics, which it does not: two training
+        scenes sweep the same space and must not share a store.
         """
         axes = {str(d): space[d].values.tolist() for d in space.dims}
         statics = {k: repr(v) for k, v in self._statics().items()}
         digest = hashlib.sha1(
-            repr((sorted(axes.items()), sorted(statics.items()))).encode()
+            repr(
+                (
+                    sorted(axes.items()),
+                    sorted(statics.items()),
+                    sorted(self._bound_hashes(bound).items()),
+                )
+            ).encode()
         ).hexdigest()
         return f"{type(self).__name__}-{digest[:8]}"
 
+    @staticmethod
+    def _bound_hashes(bound: dict[str, Any]) -> dict[str, str]:
+        """Return a stable hash of each scene object bound to the physics."""
+        hashes: dict[str, str] = {}
+        for name, value in bound.items():
+            if isinstance(value, (xr.Dataset, xr.DataArray)):
+                hashes[name] = str(joblib.hash(value))
+            else:
+                hashes[name] = repr(value)
+        return hashes
+
     def _chunks(self, space: xr.Dataset) -> dict[str, int]:
         """Return the chunk widths that give batches of ``batch_size`` states.
 

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

@@ -197,7 +197,7 @@ class TestSweepPolicy:
             remove_rayleigh=False,
             cache=CacheStore(tmp_path),
         )
-        policy = sampler._policy(sampler._space())
+        policy = sampler._policy(sampler._space(), {})
         vec = {var.name for var in sampler._contract.vec}
         assert "RhoToaSymSampler" in str(policy.store)
         assert set(policy.chunks) <= vec
@@ -225,7 +225,7 @@ class TestSweepPolicy:
             batch_size=4,
             cache=CacheStore(tmp_path),
         )
-        chunks = sampler._policy(sampler._space()).chunks
+        chunks = sampler._policy(sampler._space(), {}).chunks
         assert np.prod(list(chunks.values())) == 4
 
     def test_a_spectral_sweep_is_left_in_memory(self, tmp_path):
@@ -244,7 +244,7 @@ class TestSweepPolicy:
             remove_rayleigh=False,
             cache=CacheStore(tmp_path),
         )
-        assert sampler._policy(sampler._space()).store is UNSET
+        assert sampler._policy(sampler._space(), {}).store is UNSET
 
     def test_two_species_do_not_share_a_store(self, tmp_path):
         from adjeff.modules.samplers import RhoToaSymSampler
@@ -268,6 +268,40 @@ class TestSweepPolicy:
                 remove_rayleigh=False,
                 cache=CacheStore(tmp_path),
             )
-            return str(sampler._policy(sampler._space()).store)
+            return str(sampler._policy(sampler._space(), {}).store)
 
         assert name("ammonium") != name("blackcar")
+
+    def test_two_scenes_do_not_share_a_store(self, tmp_path):
+        from adjeff.core import S2Band, gaussian_image_dict
+        from adjeff.modules.samplers import RhoToaSymSampler
+        from adjeff.utils import CacheStore
+
+        sampler = RhoToaSymSampler(
+            atmo_config=atmo.AtmoConfig(
+                aot=xr.DataArray([0.1, 0.2], dims=["aot"]),
+                rh=xr.DataArray([50.0], dims=["rh"]),
+                h=xr.DataArray([0.0], dims=["h"]),
+                href=xr.DataArray([2.0], dims=["href"]),
+                species={"sulphate": 1.0},
+            ),
+            geo_config=atmo.GeoConfig(
+                sza=xr.DataArray([40.0], dims=["sza"]),
+                vza=xr.DataArray([0.0], dims=["vza"]),
+                saa=xr.DataArray([0.0], dims=["saa"]),
+                vaa=xr.DataArray([0.0], dims=["vaa"]),
+            ),
+            remove_rayleigh=False,
+            cache=CacheStore(tmp_path),
+        )
+        band = S2Band.B02
+        scenes = [
+            gaussian_image_dict(sigma=s, res_km=0.5, bands=[band], n=9)[band]
+            for s in (1.0, 5.0)
+        ]
+        space = sampler._space()
+        stores = [
+            str(sampler._policy(space, {"rho_s": ds, "band": band}).store)
+            for ds in scenes
+        ]
+        assert stores[0] != stores[1]