Merge pull request #56 from TUDelftGeodesy/55_from_zarr

55 from ds
TUDelftGeodesy · Oct 14, 2024 · fc04d1f · fc04d1f
2 parents b80f0ea + 6bfd60a
commit fc04d1f
Show file tree

Hide file tree

Showing 61 changed files with 845 additions and 61 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -0,0 +1,9 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+    rev: v0.6.4
+    hooks:
+      # Run the linter.
+      - id: ruff
+      # Run the formatter.
+      - id: ruff-format
diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md
@@ -35,6 +35,7 @@ to see if someone already filed the same issue;
 - (**important**) announce your plan to the rest of the community *before you start working*. This announcement should be in the form of a (new) issue;
 - (**important**) wait until some kind of consensus is reached about your idea being a good idea;
 - if needed, fork the repository to your own Github profile and create your own feature branch off of the latest master commit. While working on your feature branch, make sure to stay up to date with the master branch by pulling in changes, possibly from the 'upstream' repository (follow the instructions from GitHub: [instruction 1: configuring a remote for a fork](https://help.github.com/articles/configuring-a-remote-for-a-fork/) and [instruction 2: syncing a fork](https://help.github.com/articles/syncing-a-fork/));
+- install the pre-commit hooks by running `pre-commit install` in the project root directory;
 - make sure the existing tests still work by running, e.g. `pytest tests`;
 - add your own tests (if necessary);
 - update or expand the documentation;

diff --git a/pyproject.toml b/pyproject.toml
@@ -46,6 +46,7 @@ dev = [
     "pytest-cov",
     "pycodestyle",
     "ruff",
+    "pre-commit",
 ]
 docs = [
     "mkdocs",
@@ -89,7 +90,7 @@ ignore = [
 ] # docstring style
 
 line-length = 88
-exclude = ["docs", "build"]
+exclude = ["docs", "build", "examples"]
 # Allow unused variables when underscore-prefixed.
 dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
 target-version = "py310"

diff --git a/sarxarray/__init__.py b/sarxarray/__init__.py
@@ -1,5 +1,5 @@
 from sarxarray import stack
-from sarxarray._io import from_binary
+from sarxarray._io import from_binary, from_ds
 from sarxarray.utils import complex_coherence, multi_look
 
-__all__ = ("stack", "from_binary", "multi_look", "complex_coherence")
+__all__ = ("stack", "from_binary", "from_ds", "multi_look", "complex_coherence")
diff --git a/sarxarray/_io.py b/sarxarray/_io.py
@@ -13,6 +13,60 @@
 # Example: https://docs.dask.org/en/stable/array-creation.html#memory-mapping
 
 
+def from_ds(ds: xr.Dataset) -> xr.Dataset:
+    """Create a SLC stack or from an Xarray Dataset.
+
+    This function create tasks graph converting the two data variables of complex data:
+    `real` and `imag`, to three variables: `complex`, `amplitude`, and `phase`.
+
+    The function is intented for an SLC stack in `xr.Dataset` loaded from a Zarr file.
+
+    For other datasets, such as lat, lon, etc., please use `xr.open_zarr` directly.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        SLC stack loaded from a Zarr file.
+        Must have three dimensions: `(azimuth, range, time)`.
+        Must have two variables: `real` and `imag`.
+
+    Returns
+    -------
+    xr.Dataset
+        Converted SLC stack.
+        An xarray.Dataset with three dimensions: `(azimuth, range, time)`, and
+        three variables: `complex`, `amplitude`, `phase`.
+
+    Raises
+    ------
+    ValueError
+        The input dataset should have three dimensions: `(azimuth, range, time)`.
+    ValueError
+        The input dataset should have the following variables: `('real', 'imag')`.
+    """
+    # Check ds should have the following dimensions: (azimuth, range, time)
+    if any(dim not in ds.dims for dim in ["azimuth", "range", "time"]):
+        raise ValueError(
+            "The input dataset should have three dimensions: (azimuth, range, time)."
+        )
+
+    # Check ds should have the following variables: ("real", "imag")
+    if any(var not in ds.variables for var in ["real", "imag"]):
+        raise ValueError(
+            "The input dataset should have the following variables: ('real', 'imag')."
+        )
+
+    # Construct the three datavariables: complex, amplitude, and phase
+    ds["complex"] = ds["real"] + 1j * ds["imag"]
+    ds = ds.slcstack._get_amplitude()
+    ds = ds.slcstack._get_phase()
+
+    # Remove the original real and imag variables
+    ds = ds.drop_vars(["real", "imag"])
+
+    return ds
+
+
 def from_binary(
     slc_files, shape, vlabel="complex", dtype=np.complex64, chunks=None, ratio=1
 ):
@@ -43,11 +97,9 @@ def from_binary(
     if not np.dtype(dtype).isbuiltin:
         if not all([name in (("re", "im")) for name in dtype.names]):
             raise TypeError(
-
-                    "The customed dtype should have only two field names: "
-                    '"re" and "im". For example: '
-                    'dtype = np.dtype([("re", np.float32), ("im", np.float32)]).'
-
+                "The customed dtype should have only two field names: "
+                '"re" and "im". For example: '
+                'dtype = np.dtype([("re", np.float32), ("im", np.float32)]).'
             )
 
     # Initialize stack as a Dataset

diff --git a/tests/data/zarrs/slcs_example.zarr/.zattrs b/tests/data/zarrs/slcs_example.zarr/.zattrs
@@ -0,0 +1 @@
+{}
diff --git a/tests/data/zarrs/slcs_example.zarr/.zgroup b/tests/data/zarrs/slcs_example.zarr/.zgroup
@@ -0,0 +1,3 @@
+{
+    "zarr_format": 2
+}
diff --git a/tests/data/zarrs/slcs_example.zarr/.zmetadata b/tests/data/zarrs/slcs_example.zarr/.zmetadata
@@ -0,0 +1,146 @@
+{
+    "metadata": {
+        ".zattrs": {},
+        ".zgroup": {
+            "zarr_format": 2
+        },
+        "azimuth/.zarray": {
+            "chunks": [
+                350
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<i8",
+            "fill_value": null,
+            "filters": null,
+            "order": "C",
+            "shape": [
+                2
+            ],
+            "zarr_format": 2
+        },
+        "azimuth/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "azimuth"
+            ]
+        },
+        "imag/.zarray": {
+            "chunks": [
+                350,
+                1750,
+                1
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<f2",
+            "fill_value": "NaN",
+            "filters": null,
+            "order": "C",
+            "shape": [
+                2,
+                2,
+                3
+            ],
+            "zarr_format": 2
+        },
+        "imag/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "azimuth",
+                "range",
+                "time"
+            ]
+        },
+        "range/.zarray": {
+            "chunks": [
+                1750
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<i8",
+            "fill_value": null,
+            "filters": null,
+            "order": "C",
+            "shape": [
+                2
+            ],
+            "zarr_format": 2
+        },
+        "range/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "range"
+            ]
+        },
+        "real/.zarray": {
+            "chunks": [
+                350,
+                1750,
+                1
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<f2",
+            "fill_value": "NaN",
+            "filters": null,
+            "order": "C",
+            "shape": [
+                2,
+                2,
+                3
+            ],
+            "zarr_format": 2
+        },
+        "real/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "azimuth",
+                "range",
+                "time"
+            ]
+        },
+        "time/.zarray": {
+            "chunks": [
+                3
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<i8",
+            "fill_value": null,
+            "filters": null,
+            "order": "C",
+            "shape": [
+                3
+            ],
+            "zarr_format": 2
+        },
+        "time/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "time"
+            ]
+        }
+    },
+    "zarr_consolidated_format": 1
+}
diff --git a/tests/data/zarrs/slcs_example.zarr/azimuth/.zarray b/tests/data/zarrs/slcs_example.zarr/azimuth/.zarray
@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        350
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<i8",
+    "fill_value": null,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        2
+    ],
+    "zarr_format": 2
+}
diff --git a/tests/data/zarrs/slcs_example.zarr/azimuth/.zattrs b/tests/data/zarrs/slcs_example.zarr/azimuth/.zattrs
@@ -0,0 +1,5 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "azimuth"
+    ]
+}
diff --git a/tests/data/zarrs/slcs_example.zarr/azimuth/0 b/tests/data/zarrs/slcs_example.zarr/azimuth/0
diff --git a/tests/data/zarrs/slcs_example.zarr/imag/.zarray b/tests/data/zarrs/slcs_example.zarr/imag/.zarray
@@ -0,0 +1,24 @@
+{
+    "chunks": [
+        350,
+        1750,
+        1
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<f2",
+    "fill_value": "NaN",
+    "filters": null,
+    "order": "C",
+    "shape": [
+        2,
+        2,
+        3
+    ],
+    "zarr_format": 2
+}
diff --git a/tests/data/zarrs/slcs_example.zarr/imag/.zattrs b/tests/data/zarrs/slcs_example.zarr/imag/.zattrs
@@ -0,0 +1,7 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "azimuth",
+        "range",
+        "time"
+    ]
+}
diff --git a/tests/data/zarrs/slcs_example.zarr/imag/0.0.0 b/tests/data/zarrs/slcs_example.zarr/imag/0.0.0
diff --git a/tests/data/zarrs/slcs_example.zarr/imag/0.0.1 b/tests/data/zarrs/slcs_example.zarr/imag/0.0.1
diff --git a/tests/data/zarrs/slcs_example.zarr/imag/0.0.2 b/tests/data/zarrs/slcs_example.zarr/imag/0.0.2
diff --git a/tests/data/zarrs/slcs_example.zarr/range/.zarray b/tests/data/zarrs/slcs_example.zarr/range/.zarray
@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        1750
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<i8",
+    "fill_value": null,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        2
+    ],
+    "zarr_format": 2
+}
diff --git a/tests/data/zarrs/slcs_example.zarr/range/.zattrs b/tests/data/zarrs/slcs_example.zarr/range/.zattrs
@@ -0,0 +1,5 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "range"
+    ]
+}
diff --git a/tests/data/zarrs/slcs_example.zarr/range/0 b/tests/data/zarrs/slcs_example.zarr/range/0
diff --git a/tests/data/zarrs/slcs_example.zarr/real/.zarray b/tests/data/zarrs/slcs_example.zarr/real/.zarray
@@ -0,0 +1,24 @@
+{
+    "chunks": [
+        350,
+        1750,
+        1
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<f2",
+    "fill_value": "NaN",
+    "filters": null,
+    "order": "C",
+    "shape": [
+        2,
+        2,
+        3
+    ],
+    "zarr_format": 2
+}