Source code for torx.analysis.lineouts.lineout_set_m

"""Defines 3D Lineouts made up of collections of 2D Lineouts."""
import numpy as np
import xarray as xr
import dask.bag as db
from functools import cached_property
from typing import Union

from torx.arrays import combine_obj_list
from torx.units import check_units
from torx.decorators import AutoDocClass
from torx.grid import Grid2D, Grid3D



class LineoutSet(AutoDocClass):
    """Groups a set of 2D lineouts for a given dimension. Defaults to phi."""



    def __init__(self, lineout_array: xr.DataArray, dim_name: str = 'phi'):
        """Initialize the 3D lineout."""
        self.lineout_array = lineout_array
        self.coords = lineout_array.coords
        self.dims = lineout_array.dims
        self.attrs = lineout_array.attrs
        self.name = 'LineoutSet'
        self.encoding = lineout_array.encoding
        self.sizes = lineout_array.sizes
        self.dim_name = dim_name
        self.len = lineout_array.sizes[self.dim_name]
        self.lineout_db = db.from_sequence(
                        [lineout_array.isel(
                                            {self.dim_name: i}
                                            ).item() for i in range(self.len)],
                        npartitions=self.len
                        )


    @property
    def lineout_array(self):
        """Return the arrays of values in the grid."""
        return self._lineout_array

    @lineout_array.setter
    def lineout_array(self, value):
        """Set the arrays of values in the lineout."""
        self._lineout_array = value

    @cached_property
    def r_source(self) -> xr.DataArray:
        """Source R coordinates for the lineout (normalized to R0)."""
        return combine_obj_list(self.lineout_array.values, "r_source", 
                                self.dim_name, 
                                self.lineout_array.coords[self.dim_name].values
                                ).rename({"dim_0": "points"})

    @cached_property
    def z_source(self) -> xr.DataArray:
        """Source Z coordinates for the lineout (normalized to R0)."""
        return combine_obj_list(self.lineout_array.values, "z_source", 
                                self.dim_name, 
                                self.lineout_array.coords[self.dim_name].values
                                ).rename({"dim_0": "points"})

    @cached_property
    def r_points(self) -> xr.DataArray:
        """R coordinates for the lineout (normalized to R0)."""
        try:
            return combine_obj_list(self.lineout_array.values, "r_points", 
                                    self.dim_name, 
                                    self.lineout_array.coords[self.dim_name].values
                                    ).rename({"dim_0": "points"})
        except AttributeError as e:
            msg = "r_points are not defined. "+\
                  "You likely need to run find_points_on_grid first."
            raise AttributeError(msg) from e

    @cached_property
    def z_points(self) -> xr.DataArray:
        """Z coordinates for the lineout (normalized to R0)."""
        try:
            return combine_obj_list(self.lineout_array.values, "z_points", 
                                    self.dim_name, 
                                    self.lineout_array.coords[self.dim_name].values
                                    ).rename({"dim_0": "points"})
        except AttributeError as e:
            msg = "z_points are not defined. "+\
                  "You likely need to run find_points_on_grid first."
            raise AttributeError(msg) from e

    @cached_property
    def is_periodic(self):
        """Determine whether all lineouts are periodic, i.e. a closed polygon."""
        return all(
                  self.lineout_db.map(lambda lineout: lineout.is_periodic
                  ).compute())

    def find_points_on_grid(self, grid: Union[Grid2D, Grid3D], n_samples: int=500):
        """Run find_points_on_grid for each lineout in parallel using Dask."""
        self.n_samples = n_samples

        if isinstance(grid, Grid3D):
            self.is_3D = True
            grid_db = db.from_sequence(
                [
                grid.grid_array.isel({self.dim_name: i}).item() \
                for i in range(len(grid.grid_array))
                ],
                npartitions=self.len
                )
            # Zip delivers a (lineout, grid_slice) tuple to the helper
            updated_lineouts = db.zip(
                                      self.lineout_db, 
                                      grid_db
                                      ).map(self._find_points_helper).compute()
            
        elif isinstance(grid, Grid2D):
            self.is_3D = False
            # Pass the single grid as a keyword argument to the helper
            updated_lineouts = self.lineout_db.map(self._find_points_helper, 
                                                   grid_2d=grid).compute()
        else:
            raise TypeError(f"Expected Grid3D or Grid2D, got {type(grid)}")

        # Update the saved lineout xarray and dask bag
        self.lineout_array.values = np.array(updated_lineouts)
        self.lineout_db = db.from_sequence(
            self.lineout_array.values,
            npartitions=self.len
        )

    def interpolate(self, input_array: xr.DataArray) -> xr.DataArray:
        """Interpolates the input array using each lineout for a given dimension."""
        if not hasattr(self, 'is_3D'):
            msg = "is_3D is not defined. Likely must run find_points_on_grid."
            raise AttributeError(msg)
        if self.is_3D:
            input_array_db = db.from_sequence(
                [input_array.isel({self.dim_name: i}) for i in range(self.len)],
                npartitions=self.len
                )
            zipped_bag = db.zip(self.lineout_db, input_array_db)
            interpolated_bag = zipped_bag.map(
                lambda p: p[0].interpolate(
                    p[1].isel(points=slice(None, p[0].grid_size))
                )
            ).compute()
        else:
            interpolated_bag = self.lineout_db.map(
                               lambda lineout: lineout.interpolate(input_array)
                               ).compute()
        
        return xr.concat(
            interpolated_bag, 
            dim=self.dim_name
        ).assign_coords({self.dim_name: self.lineout_array.coords[self.dim_name]})

    def _find_points_helper(self, data, grid_2d=None):
            """Call find_points_on_grid for a single lineout on a 2D or 3D grid."""
            # If it's a tuple (from db.zip in 3D), unpack it
            if isinstance(data, tuple):
                lineout, grid = data
            else:
                # If it's just the lineout (2D), use the static grid provided
                lineout = data
                grid = grid_2d
                
            lineout.find_points_on_grid(grid, n_samples=self.n_samples)
            return lineout