From 6bfd60ab7ee720f315513f9c0b486e0630b24085 Mon Sep 17 00:00:00 2001
From: Ou Ku <o.ku@esciencecenter.nl>
Date: Fri, 11 Oct 2024 11:32:09 +0200
Subject: [PATCH] rename from_zarr to from_ds

---
 sarxarray/__init__.py |  4 ++--
 sarxarray/_io.py      | 29 ++++++++++++++---------------
 tests/test_io.py      | 28 ++++++++++++++++------------
 3 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/sarxarray/__init__.py b/sarxarray/__init__.py
index 54c821c..d4393db 100644
--- a/sarxarray/__init__.py
+++ b/sarxarray/__init__.py
@@ -1,5 +1,5 @@
 from sarxarray import stack
-from sarxarray._io import from_binary, from_zarr
+from sarxarray._io import from_binary, from_ds
 from sarxarray.utils import complex_coherence, multi_look
 
-__all__ = ("stack", "from_binary", "from_zarr", "multi_look", "complex_coherence")
+__all__ = ("stack", "from_binary", "from_ds", "multi_look", "complex_coherence")
diff --git a/sarxarray/_io.py b/sarxarray/_io.py
index fc9d4b3..86efe60 100644
--- a/sarxarray/_io.py
+++ b/sarxarray/_io.py
@@ -1,6 +1,5 @@
 import logging
 import math
-from pathlib import Path
 
 import dask
 import dask.array as da
@@ -14,37 +13,37 @@
 # Example: https://docs.dask.org/en/stable/array-creation.html#memory-mapping
 
 
-def from_zarr(files_zarr: str | Path) -> xr.Dataset:
-    """Read a SLC stack or related variables from zarr files.
+def from_ds(ds: xr.Dataset) -> xr.Dataset:
+    """Create a SLC stack or from an Xarray Dataset.
 
-    Note that this function only works for the SLC stack zarr files.
+    This function create tasks graph converting the two data variables of complex data:
+    `real` and `imag`, to three variables: `complex`, `amplitude`, and `phase`.
 
-    The purpose of this function is to create task graph converting `real`/`imag` to
-    `complex`, `amplitude`, and `phase` variables.
+    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
     ----------
-    files_zarr : str | Path
-        Paths to the zarr files.
+    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
-        Loaded SLC stack.
-        An xarray.Dataset with three dimensions: (azimuth, range, time), and
-        three variables: ("complex", "amplitude", "phase").
+        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).
+        The input dataset should have three dimensions: `(azimuth, range, time)`.
     ValueError
-        The input dataset should have the following variables: ('real', 'imag').
+        The input dataset should have the following variables: `('real', 'imag')`.
     """
-    ds = xr.open_zarr(files_zarr)
-
     # 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(
diff --git a/tests/test_io.py b/tests/test_io.py
index 93b2751..f570d17 100644
--- a/tests/test_io.py
+++ b/tests/test_io.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 import pytest
+import xarray as xr
 
 import sarxarray
 from sarxarray._io import _calc_chunksize, _unpack_complex
@@ -17,30 +18,33 @@ def test_slcs():
     ]
 
 
-class TestFromZarr:
-    """from_zarr in _io.py"""
+class TestFromDS:
+    """from_ds in _io.py"""
 
-    def test_from_zarr_normal(self):
-        slcs = sarxarray.from_zarr(
+    def test_from_ds_normal(self):
+        test_ds = xr.open_zarr(
             f"{os.path.dirname(__file__)}/data/zarrs/slcs_example.zarr"
         )
+        slcs = sarxarray.from_ds(test_ds)
         assert all(dim in slcs.dims for dim in ["azimuth", "range", "time"])
         assert all(var not in slcs.variables.keys() for var in ["real", "imag"])
         assert all(
             var in slcs.variables.keys() for var in ["complex", "amplitude", "phase"]
         )
 
-    def test_from_zarr_broken_dim(self):
+    def test_from_ds_broken_dim(self):
+        test_ds_broken_dim = xr.open_zarr(
+            f"{os.path.dirname(__file__)}/data/zarrs/slcs_example_broken_dim.zarr"
+        )
         with pytest.raises(ValueError):
-            sarxarray.from_zarr(
-                f"{os.path.dirname(__file__)}/data/zarrs/slcs_example_broken_dim.zarr"
-            )
+            sarxarray.from_ds(test_ds_broken_dim)
 
-    def test_from_zarr_broken_vars(self):
+    def test_from_ds_broken_vars(self):
+        test_ds_broken_vars = xr.open_zarr(
+            f"{os.path.dirname(__file__)}/data/zarrs/slcs_example_broken_vars.zarr"
+        )
         with pytest.raises(ValueError):
-            sarxarray.from_zarr(
-                f"{os.path.dirname(__file__)}/data/zarrs/slcs_example_broken_vars.zarr"
-            )
+            sarxarray.from_ds(test_ds_broken_vars)
 
 
 class TestFromBinary: