Source code for zefir.potential.flows.source

"""
Source/Sink flow implementation.

This module provides the SourceFlow class representing a source or sink flow
with radial flow pattern emanating from or converging to a point.
"""

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




class SourceFlow(ComplexPotential):
    """A source/sink flow defined by its flow rate and center.
    
    A source (or sink) flow represents fluid emanating from (or converging to)
    a point. The complex potential is given by:
    
    .. math:: f(z) = \\frac{q_v}{2\\pi} \\log(z - z_0)
    
    where :math:`q_v` is the flow rate (positive for source, negative for sink)
    and :math:`z_0` is the center position.
    
    Parameters
    ----------
    flow_rate : float
        Flow rate q_v (positive for source, negative for sink)
    center : Tuple[float, float] or complex
        Position (x, y) or complex number representing the center
        
    Attributes
    ----------
    _center : complex
        Complex representation of the center position
    _intensity : float
        Scaled intensity (flow_rate / 2π)
    _original_flow_rate : float
        Original flow rate value
        
    Examples
    --------
    >>> flow = SourceFlow(1.0, (0.5, 0.0))
    >>> z = np.array([1+0j, 2+0j, 0+1j])
    >>> potential = flow(z)
    >>> velocity = flow.velocity(z)
    """
    
    metadata = FlowMetadata(
        "Source/Sink",
        {
            "flow_rate": {
                "type": ParamType.FLOAT,
                "default": 1.0,
                "min": -100,
                "max": 100
            },
            "center": {
                "type": ParamType.VECTOR2,
                "default": [0.0, 0.0],
                "min": -10,
                "max": 10,
                "labels": ["X center", "Y center"]
            }
        },
        has_position=True
    )
    


    def __init__(self, flow_rate:float, center:Tuple[float, float]|complex):
        if isinstance(center, complex):
            self._center = center
        else:
            self._center = complex(*center)
        self._intensity = flow_rate/(2*np.pi)
        self._original_flow_rate = flow_rate




    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._intensity*np.log(z-self._center)
        elif deriv == 1:
            return self._intensity/(z-self._center)
        else: 
            raise InvalidDerivLevel()

    


    def update_parameters(self, params: dict):
        """Update flow parameters from a dictionary.
        
        Parameters
        ----------
        params : dict
            Dictionary with keys "flow_rate" and/or "center"
        """
        for param_name, value in params.items():
            if param_name == "flow_rate":
                self._intensity = value / (2 * np.pi)
                self._original_flow_rate = value
            elif param_name == "center":
                if isinstance(value, (list, tuple)):
                    value = complex(*value)
                self._center = value
            elif hasattr(self, f'_{param_name}'):
                setattr(self, f'_{param_name}', value)

    


    def get_parameters(self) -> dict:
        """Get current parameter values as a dictionary.
        
        Returns
        -------
        dict
            Dictionary with "flow_rate" and "center" keys
        """
        return {
            "flow_rate": self._original_flow_rate,
            "center": [self._center.real, self._center.imag]
        }