Source code for torx.analysis.lineouts.lineout_set_m

"""Defines 3D Lineouts made up of collections of 2D Lineouts."""
from abc import ABC
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 autodoc_class
from torx.grid import Grid2D, Grid3D



@autodoc_class
class LineoutSet(ABC):
    """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