From bb076a26570ee0d398d210e7afec1ab6d8209f53 Mon Sep 17 00:00:00 2001 From: Ryuichi Arafune Date: Fri, 20 Oct 2023 12:32:42 +0900 Subject: [PATCH 1/2] =?UTF-8?q?=F0=9F=9A=A7=20=20test=20for=20G.stride=20a?= =?UTF-8?q?nd=20G.range?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/test_xarray_extensions.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/test_xarray_extensions.py b/tests/test_xarray_extensions.py index faf64969..3f985052 100644 --- a/tests/test_xarray_extensions.py +++ b/tests/test_xarray_extensions.py @@ -89,6 +89,27 @@ def test_full_coords(self, dataset_cut: xr.Dataset) -> None: "eV", ] + def test_G_stride_and_range(self, dataarray_cut: xr.DataArray) -> None: + """Test for G.range and G.stride.""" + generic_stride = dataarray_cut.G.stride(generic_dim_names=True) + assert "x" in generic_stride + assert "y" in generic_stride + generic_stride = dataarray_cut.G.stride("eV", generic_dim_names=False) + np.testing.assert_almost_equal(generic_stride, 0.0023255810) + stride = dataarray_cut.G.stride(["eV"], generic_dim_names=False) + np.testing.assert_almost_equal(stride, 0.0023255810) + stride = dataarray_cut.G.stride(["eV", "phi"], generic_dim_names=False) + np.testing.assert_array_almost_equal(stride, (0.0023255810, 0.001745)) + # + range_ = dataarray_cut.G.range(generic_dim_names=False) + np.testing.assert_array_almost_equal(range_["eV"], (-0.4255814, 0.13023245)) + np.testing.assert_array_almost_equal( + range_["phi"], (0.22165681500327986, 0.6387905062299246) + ) + range_ = dataarray_cut.G.range(generic_dim_names=True) + np.testing.assert_array_almost_equal(range_["y"], (-0.4255814, 0.13023245)) + np.testing.assert_array_almost_equal(range_["x"], (0.22165681500327986, 0.6387905062299246)) + def test_experimental_conditions(self, dataset_cut: xr.Dataset) -> None: """Test for property experimenta_conditions.""" assert dataset_cut.S.experimental_conditions == { From 283346a4b687084714ebdec06f69bffc9e60b009 Mon Sep 17 00:00:00 2001 From: Ryuichi Arafune Date: Fri, 20 Oct 2023 15:22:57 +0900 Subject: [PATCH 2/2] =?UTF-8?q?=F0=9F=92=AC=20=20update=20type=20hints?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- arpes/analysis/path.py | 13 ++++++------ arpes/analysis/pocket.py | 35 ++++++++++++++++++--------------- tests/test_xarray_extensions.py | 3 ++- 3 files changed, 27 insertions(+), 24 deletions(-) diff --git a/arpes/analysis/path.py b/arpes/analysis/path.py index baad16af..bc4366b2 100644 --- a/arpes/analysis/path.py +++ b/arpes/analysis/path.py @@ -98,14 +98,14 @@ def to_dataarray(name: str) -> xr.DataArray: @update_provenance("Select from data along a path") def select_along_path( path: xr.Dataset, - data: DataType, + data: xr.DataArray, radius: float = 0, n_points: int = 0, *, fast: bool = True, scaling: float | xr.Dataset | dict[str, NDArray[np.float_]] | None = None, **kwargs: Incomplete, -) -> DataType: +) -> xr.DataArray: """Performs integration along a path. This functionally allows for performing a finite width @@ -117,11 +117,10 @@ def select_along_path( path: The path to select along. data: The data to select/interpolate from. radius: A number or dictionary of radii to use for the selection along different dimensions, - if none is provided reasonable values will be chosen. Alternatively, you can pass radii - via `{dim}_r` kwargs as well, i.e. 'eV_r' or 'kp_r' + if none is provided reasonable values will be chosen. Alternatively, you can pass + radii via `{dim}_r` kwargs as well, i.e. 'eV_r' or 'kp_r' n_points: The number of points to interpolate along the path, by default we will infer a - reasonable number from the radius parameter, if provided or inferred - fast: If fast is true, will use rectangular selections rather than ellipsoid ones + reasonable number from the radius parameter, if provided or inferred Returns: The data selected along the path. @@ -130,6 +129,6 @@ def select_along_path( selections = [] for _, view in new_path.G.iterate_axis("index"): - selections.append(data.S.select_around(view, radius=radius, fast=fast, **kwargs)) + selections.append(data.S.select_around(view, radius=radius, **kwargs)) return xr.concat(selections, new_path.index) diff --git a/arpes/analysis/pocket.py b/arpes/analysis/pocket.py index b2eb09de..e121b52e 100644 --- a/arpes/analysis/pocket.py +++ b/arpes/analysis/pocket.py @@ -146,10 +146,7 @@ def radial_edcs_along_pocket( selection_coords = [{k: v[n] for k, v in data_vars.items()} for n in range(n_points)] - edcs = [ - data_array.S.select_around(coord, radius=select_radius, fast=True) - for coord in selection_coords - ] + edcs = [data_array.S.select_around(coord, radius=select_radius) for coord in selection_coords] for r, edc in zip(radius_coord, edcs, strict=True): edc.coords["r"] = r @@ -161,10 +158,9 @@ def radial_edcs_along_pocket( def curves_along_pocket( data: DataType, - n_points=None, - inner_radius: float = 0, - outer_radius: float = 5, - shape=None, + n_points: int = 0, + inner_radius: float = 0.0, + outer_radius: float = 5.0, **kwargs: Incomplete, ) -> tuple[list[xr.DataArray], list[float]]: """Produces radial slices along a Fermi surface through a pocket. @@ -178,7 +174,7 @@ def curves_along_pocket( Args: data: input data - n_points: + n_points (int): inner_radius: outer_radius: shape: @@ -197,7 +193,7 @@ def curves_along_pocket( center_as_vector = np.array([center_point.get(d, 0) for d in fermi_surface_dims]) - if n_points is None: + if not n_points: # determine N approximately by the granularity n_points = np.min([len(data_array.coords[d].values) for d in fermi_surface_dims]) @@ -206,13 +202,16 @@ def curves_along_pocket( angles = np.linspace(0, 2 * np.pi, n_points, endpoint=False) - def slice_at_angle(theta: float): + def slice_at_angle(theta: float) -> xr.DataArray: primitive = np.array([np.cos(theta), np.sin(theta)]) far = center_as_vector + outer_radius * primitive return slice_along_path( data_array, - [dict(zip(fermi_surface_dims, point)) for point in [center_as_vector, far]], + [ + dict(zip(fermi_surface_dims, point, strict=True)) + for point in [center_as_vector, far] + ], resolution=resolution, ) @@ -235,14 +234,18 @@ def slice_at_angle(theta: float): return slices, angles -def find_kf_by_mdc(slice_data: DataType, offset: float = 0, **kwargs: Incomplete) -> float: +def find_kf_by_mdc( + slice_data: DataType, + offset: float = 0, + **kwargs: Incomplete, +) -> float: """Finds the Fermi momentum by curve fitting an MDC. Offset is used to control the radial offset from the pocket for studies where you want to go slightly off the Fermi momentum. Args: - slice: Input fit data. + slice_data: Input fit data. offset: Offset to add to the result kwargs: Passed as parameters to the fit routine. @@ -268,7 +271,7 @@ def find_kf_by_mdc(slice_data: DataType, offset: float = 0, **kwargs: Incomplete def edcs_along_pocket( data: DataType, kf_method: Callable[..., float] | None = None, - select_radius=None, + select_radius: dict[str, float] | None = None, sel: dict[str, slice] | None = None, method_kwargs: Incomplete | None = None, **kwargs: Incomplete, @@ -308,7 +311,7 @@ def edcs_along_pocket( for kf, ss in zip(kfs, slices) ] - edcs = [data.S.select_around(_, radius=select_radius, fast=True) for _ in locations] + edcs = [data.S.select_around(_, radius=select_radius) for _ in locations] data_vars = {} index = np.array(angles) diff --git a/tests/test_xarray_extensions.py b/tests/test_xarray_extensions.py index 3f985052..3ae50356 100644 --- a/tests/test_xarray_extensions.py +++ b/tests/test_xarray_extensions.py @@ -104,7 +104,8 @@ def test_G_stride_and_range(self, dataarray_cut: xr.DataArray) -> None: range_ = dataarray_cut.G.range(generic_dim_names=False) np.testing.assert_array_almost_equal(range_["eV"], (-0.4255814, 0.13023245)) np.testing.assert_array_almost_equal( - range_["phi"], (0.22165681500327986, 0.6387905062299246) + range_["phi"], + (0.22165681500327986, 0.6387905062299246), ) range_ = dataarray_cut.G.range(generic_dim_names=True) np.testing.assert_array_almost_equal(range_["y"], (-0.4255814, 0.13023245))