Source code for materials.ideal

"""Ideal Material

This module contains the IdealMaterial class, which represents an ideal
material with a fixed refractive index and extinction coefficient for all
wavelengths.

Kramer Harrison, 2024
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import optiland.backend as be
from optiland.materials.base import BaseMaterial

if TYPE_CHECKING:
    from optiland.propagation.base import BasePropagationModel




class IdealMaterial(BaseMaterial):
    """Represents an ideal material with a fixed refractive index and extinction
    coefficient for all wavelengths.

    Evaluation preserves each stored parameter's dtype when the backend's default
    precision changes, while converting to the active backend and device.

    Attributes:
        index (float): The refractive index of the material.
        absorp (float): The extinction coefficient of the material.

    """

    def __init__(
        self,
        n: float,
        k: float = 0,
        propagation_model: BasePropagationModel | None = None,
    ):
        super().__init__(propagation_model)
        self.index = be.array([n])
        self.absorp = be.array([k])

    @property
    def display_name(self) -> str:
        """Show a constant index; identify air only when extinction is zero."""
        if self.index.item() == 1 and self.absorp.item() == 0:
            return "Air"
        return str(self.index.item())

    def _cache_state(self) -> tuple | None:
        """Track live index/extinction values, including in-place writes."""
        return self._state_key((self.index, self.absorp))

    def _calculate_n(self, wavelength, **kwargs):
        """Returns the refractive index of the material.

        Args:
            wavelength (float or be.ndarray): The wavelength(s) of light in microns.
                This argument is not used by this material model as the index is
                constant.

        Returns:
            float or be.ndarray: The refractive index of the material. Returns a
            scalar if wavelength is scalar, otherwise an array of the same shape
            as wavelength, filled with the constant refractive index.
        """
        index = self._as_backend_array(self.index, preserve_dtype=True)[0]
        if be.is_array_like(wavelength) and be.size(wavelength) > 1:
            # Broadcast the parameter itself: multiplying by default-precision
            # ones can downcast a scalar parameter and lose its precision.
            return self._broadcast_like(index, wavelength)
        return index

    def _calculate_k(self, wavelength, **kwargs):
        """Returns the extinction coefficient of the material.

        Args:
            wavelength (float or be.ndarray): The wavelength(s) of light in microns.
                This argument is not used by this material model as the value is
                constant.

        Returns:
            float or be.ndarray: The extinction coefficient of the material. Returns a
            scalar if wavelength is scalar, otherwise an array of the same shape
            as wavelength, filled with the constant extinction coefficient.
        """
        absorp = self._as_backend_array(self.absorp, preserve_dtype=True)[0]
        if be.is_array_like(wavelength) and be.size(wavelength) > 1:
            # A broadcast view preserves both parameter precision and gradients.
            return self._broadcast_like(absorp, wavelength)
        return absorp



    def to_dict(self):
        """Returns a dictionary representation of the material.

        Returns:
            dict: A dictionary representation of the material.

        """
        material_dict = super().to_dict()
        material_dict.update({"index": self.index.item(), "absorp": self.absorp.item()})
        return material_dict




    @classmethod
    def from_dict(cls, data):
        """Creates a material from a dictionary representation.

        Args:
            data (dict): The dictionary representation of the material.

        Returns:
            IdealMaterial: The material.

        """
        return cls(data["index"], data.get("absorp", 0))