Source code for zefir.potential.flows.power

"""
Power law flow implementation.

This module provides the PowerFlow class representing flows of the form αz^β,
which can model corner flows and other power-law potential flows.
"""

import numpy as np
from ...typing import CplxNDArray, ParamType
from .base import ComplexPotential, FlowMetadata, InvalidDerivLevel




class PowerFlow(ComplexPotential):
    """A potential flow of the form αz^β.
    
    The power law flow represents a family of potential flows defined by:
    
    .. math:: f(z) = \\alpha z^\\beta
    
    Special cases include:
    
    - β = 1: Uniform flow (when α is real)
    - β = 2: Flow around a 90° corner
    - β = n/2: Flow around a corner with angle π/n
    
    Parameters
    ----------
    alpha : float, optional
        Coefficient α (default: 1.0)
    beta : float, optional
        Exponent β (default: 1.0, must be > 0)
        
    Attributes
    ----------
    _alpha : float
        Coefficient α
    _beta : float
        Exponent β
        
    Examples
    --------
    >>> # Uniform flow
    >>> flow1 = PowerFlow(alpha=1.0, beta=1.0)
    >>> # 90 degree corner flow
    >>> flow2 = PowerFlow(alpha=1.0, beta=2.0)
    >>> z = np.array([1+0j, 1+1j, 2+0j])
    >>> potential = flow2(z)
    """
    
    metadata = FlowMetadata(
        "Power Law",
        {
            "alpha": {
                "type": ParamType.FLOAT,
                "default": 1.0,
                "min": -10,
                "max": 10
            },
            "beta": {
                "type": ParamType.FLOAT,
                "default": 1.0,
                "min": 0.1,
                "max": 10
            }
        },
        has_position=False
    )
    


    def __init__(self, alpha:float=1.0, beta:float=1.0):
        self._alpha = alpha
        self._beta = beta




    def __call__(
            self, z: CplxNDArray, deriv:int=0
    ) -> CplxNDArray:
        """Evaluate the complex potential or its derivative.
        
        Parameters
        ----------
        z : CplxNDArray
            Complex array of points where the potential is evaluated
        deriv : int, optional
            Derivative level: 0 for potential, 1 for velocity
            
        Returns
        -------
        CplxNDArray
            Complex potential (deriv=0) or velocity (deriv=1)
            
        Raises
        ------
        InvalidDerivLevel
            If deriv is not 0 or 1
        """
        if deriv == 0:
            return self._alpha*z**self._beta
        if deriv == 1:
            return self._alpha*self._beta*z**(self._beta-1)
        else: 
            raise InvalidDerivLevel()