diff --git a/src/earthkit/data/field/component/geography.py b/src/earthkit/data/field/component/geography.py index 21e2bf7c5..7d0e48aea 100644 --- a/src/earthkit/data/field/component/geography.py +++ b/src/earthkit/data/field/component/geography.py @@ -659,6 +659,19 @@ def area(self) -> tuple: def grid_type(self) -> str | None: return None + def set(self, *args, shape_hint=None, **kwargs) -> "GeographyBase": + """Return a new GeographyBase object with updated values.""" + kwargs = self._normalise_set_kwargs(*args, **kwargs) + keys = set(kwargs.keys()) + + if keys == {"shape"}: + return self.from_dict(**kwargs) + + if keys == {"grid_spec"} or keys == {"latitudes", "longitudes"}: + return super().from_dict(kwargs, shape_hint=shape_hint) + + raise ValueError(f"Invalid {keys=} for Geography specification") + @classmethod def from_dict(cls, d) -> "EmptyGeography": if not isinstance(d, dict): diff --git a/tests/list_of_dicts/test_lod_geography.py b/tests/list_of_dicts/test_lod_geography.py index 965ed3c38..d2aa15910 100644 --- a/tests/list_of_dicts/test_lod_geography.py +++ b/tests/list_of_dicts/test_lod_geography.py @@ -153,6 +153,34 @@ def test_lod_no_latlon( assert ds[0].geography.latlons() == (None, None) +def test_lod_geo_set_on_nongeo_1(): + from earthkit.data import Field + + f = Field.from_dict({"parameter": {"variable": "t", "units": "K"}, "values": np.array([1.0, 2.0, 3.0])}) + + lats = np.array([10.0, 20.0, 30.0]) + lons = np.array([100.0, 110.0, 120.0]) + + f_new = f.set({"geography.latitudes": lats, "geography.longitudes": lons}) + + assert np.allclose(f_new.geography.latitudes(), lats) + assert np.allclose(f_new.geography.longitudes(), lons) + + +def test_lod_geo_set_on_nongeo_2(): + from earthkit.data import Field + + f = Field.from_dict({"parameter": {"variable": "t", "units": "K"}, "values": np.array([1.0, 2.0, 3.0])}) + + lats = np.array([10.0, 20.0, 30.0]) + lons = np.array([100.0, 110.0, 120.0]) + + f_new = f.set(geography={"latitudes": np.array([10.0, 20.0, 30.0]), "longitudes": np.array([100.0, 110.0, 120.0])}) + + assert np.allclose(f_new.geography.latitudes(), lats) + assert np.allclose(f_new.geography.longitudes(), lons) + + if __name__ == "__main__": from earthkit.data.utils.testing import main