Source code for xgc_analysis.plane_data

import numpy as np
from matplotlib.tri import LinearTriInterpolator
import matplotlib.colors as mcolors          #  << NEW import
from xgc_analysis.plane import Plane



class PlaneData:
    """
    Stores field data for a single plane.
    
    For a scalar field, data is a 1D NumPy array of shape (n,).
    For a vector field, data is a 2D NumPy array of shape (n, n_components), where for each vertex
    the vector components are stored contiguously in memory.
    
    Attributes:
      - n: number of vertices in this plane.
      - field_type: 'scalar' or 'vector'.
      - n_components: number of components (1 for scalar, 2 or 3 for vector fields).
      - dtype: data type for the array.
      - data: the NumPy array storing the field data.
    """
    def __init__(self, plane, data_array=None, n_components=1, dtype=np.float64):
        self.plane = plane  # geometry REFERENCE (!)
        self.n = plane.n_n
        self.n_components = n_components
        self.dtype = dtype

        if (n_components > 1):
            self.field_type = 'vector'
        else:
            self.field_type = 'scalar'


        # Pre‑allocate placeholders for flux‑surface operations so they always exist.
        # Scalar → 1‑D arrays, vector → 2‑D with n_components columns.
        nsurf = plane.nsurf
        if n_components == 1:
            self.flux_avg_psi       = np.empty(nsurf, dtype=dtype)
            self.flux_avg_plane     = np.empty(self.n, dtype=dtype)
            # `flux_avg_norm` only meaningful for vector data; skip.
        else:
            self.flux_avg_psi       = np.empty((nsurf, n_components), dtype=dtype)
            self.flux_avg_plane     = np.empty((self.n, n_components), dtype=dtype)
            self.flux_avg_norm_psi  = np.empty(nsurf, dtype=dtype)

        if (data_array is not None):
            if (self.field_type == 'vector'):
                if ((self.n != data_array.shape[0]) or (self.n_components != data_array.shape[1])):
                    raise ValueError("Field type is set to vector, but the dimensions of the input data are inconsistent")
            elif (self.field_type == 'scalar'):
                if (self.n != data_array.shape[0]):
                    raise ValueError("Field type is set to scalar, but the dimentions of the input data are inconsistent")

        self.data = np.zeros((self.n,self.n_components),dtype=self.dtype)
        if (data_array is not None):
            self.data = data_array




    def get_data(self):
        """Returns the underlying NumPy array."""
        return self.data

    


    def extract_component(self, component_index):
        """
        For vector fields only.
        
        Returns a new PlaneData instance (a scalar field) containing the specified component from the vector.
        """
        if self.field_type != 'vector':
            raise ValueError("extract_component is only applicable for vector fields")
        if component_index < 0 or component_index >= self.n_components:
            raise IndexError("Component index out of range")
        new_pd = PlaneData(self.n, field_type='scalar', n_components=1, dtype=self.dtype)
        new_pd.data = self.data[:, component_index].copy()
        return new_pd


    # --- Flux-surface averaging ---


    def flux_avg_to_surf(self):
        """Compute flux-surface-average and store it in `self.flux_avg_psi`."""
        self.flux_avg_psi[...] = self.plane.flux_avg_to_surf(self.data)




    def flux_avg_from_surf(self):
        """Back‑project `flux_avg_psi` and store result in `self.flux_avg_plane`."""
        self.flux_avg_plane[...] = self.plane.flux_avg_from_surf(self.flux_avg_psi)




    def flux_avg_to_from_surf(self):
        """Compute the flux-surface average of `self.data` and project back to the plane in `self.flux_avg_plane`."""
        self.flux_avg_plane[...] = self.plane.flux_avg_to_from_surf(self.data)




    def flux_avg_to_surf_norm(self):
        """Compute flux-surface-average of the norm of a vector field and store it in `self.flux_avg_norm_psi`."""
        if self.field_type != 'vector':
            raise ValueError("flux_avg_to_surf_norm is only applicable for vector fields.")
        self.flux_avg_norm_psi[...] = self.plane.flux_avg_to_surf_norm(self.data)


    # --- Gradients ---


    def evaluate_scalar_gradients(self) -> "PlaneData":
        """
        Compute ∇_{R,Z}f and return the result as a new PlaneData instance
        (two-component vector on the same plane).

        Returns
        -------
        PlaneData            # field_type = 'vector'
        """
        if getattr(self, "field_type", "scalar") != "scalar":
            raise TypeError("Gradient evaluation is only defined for scalar PlaneData.")

        grad_arr = self.plane.evaluate_scalar_gradients(self.data)

        # New PlaneData on this plane, now holding a 2-component vector field
        return PlaneData(
            plane=self.plane,
            data=grad_arr,
            field_type="vector"     # ← use whatever keyword your class expects
        )


    # --- Plotting ---


    def plot_contour(
            self,
            plane: Plane,
            **kwargs,
    ):
        """
        Convenience wrapper that forwards to Plane.plot_contour using this
        instance’s data.

        Any keyword arguments accepted by Plane.plot_contour can be supplied.
        """
        plane.plot_contour(self.data, **kwargs)