Vector Fields

class VectorField(VMobject):
    def __init__(
        self,
        func: Callable[[VectArray], VectArray],
        coordinate_system: CoordinateSystem,
        density: float = 2.0,
        magnitude_range: Optional[Tuple[float, float]] = None,
        color: Optional[ManimColor] = None,
        color_map_name: Optional[str] = "3b1b_colormap",
        color_map: Optional[Callable[[Sequence[float]], Vect4Array]] = None,
        stroke_opacity: float = 1.0,
        stroke_width: float = 3,
        tip_width_ratio: float = 4,
        tip_len_to_width: float = 0.01,
        max_vect_len: float | None = None,
        max_vect_len_to_step_size: float = 0.8,
        flat_stroke: bool = False,
        norm_to_opacity_func=None,
        **kwargs
    ):
        """
        Create a visual vector field representation.
        
        Parameters:
        - func: Vectorized function mapping coordinates to vectors
        - coordinate_system: Axes or NumberPlane for coordinate mapping
        - density: Sampling density for vector placement
        - magnitude_range: Optional range for magnitude scaling
        - color_map_name: Built-in color map for magnitude coloring
        - stroke_width: Arrow thickness
        - tip_width_ratio: Arrow head width relative to shaft
        - max_vect_len: Maximum display length for vectors
        - max_vect_len_to_step_size: Scaling factor for vector length
        """
    
    def update_vectors(self):
        """Recompute and redraw all vector arrows based on function."""
    
    def set_sample_coords(self, sample_coords):
        """Update the grid of sample points for vector placement."""
    
    def set_stroke_width(self, width):
        """Adjust arrow thickness with proper tip scaling."""

class TimeVaryingVectorField(VectorField):
    def __init__(
        self,
        time_func: Callable[[VectArray, float], VectArray],
        coordinate_system: CoordinateSystem,
        **kwargs
    ):
        """
        Create a time-dependent vector field.
        
        Parameters:
        - time_func: Function taking (coordinates, time) returning vectors
        - coordinate_system: Coordinate system for field mapping
        - kwargs: Additional VectorField parameters
        """
    
    def increment_time(self, dt):
        """Advance the internal time counter by dt."""

class StreamLines(VGroup):
    def __init__(
        self,
        func: Callable[[VectArray], VectArray],
        coordinate_system: CoordinateSystem,
        density: float = 1.0,
        n_repeats: int = 1,
        noise_factor: float | None = None,
        solution_time: float = 3,
        dt: float = 0.05,
        arc_len: float = 3,
        max_time_steps: int = 200,
        n_samples_per_line: int = 10,
        cutoff_norm: float = 15,
        stroke_width: float = 1.0,
        stroke_color: ManimColor = WHITE,
        stroke_opacity: float = 1,
        color_by_magnitude: bool = True,
        magnitude_range: Tuple[float, float] = (0, 2.0),
        taper_stroke_width: bool = False,
        color_map: str = "3b1b_colormap",
        **kwargs
    ):
        """
        Generate streamlines following vector field flow.
        
        Parameters:
        - func: Vector field function
        - coordinate_system: Coordinate system for integration
        - density: Sampling density for streamline starting points
        - n_repeats: Multiple lines per sample point
        - noise_factor: Random offset for varied starting points
        - solution_time: Integration time duration
        - dt: Time step for numerical integration
        - max_time_steps: Maximum integration steps
        - cutoff_norm: Stop integration if field norm exceeds this
        - color_by_magnitude: Color lines by local field strength
        - taper_stroke_width: Varying line thickness along streamline
        """
    
    def draw_lines(self):
        """Use ODE solver to generate streamline paths."""
    
    def get_sample_coords(self):
        """Get starting points with optional noise."""

class AnimatedStreamLines(VGroup):
    def __init__(
        self,
        stream_lines: StreamLines,
        lag_range: float = 4,
        rate_multiple: float = 1.0,
        line_anim_config: dict = {"rate_func": linear, "time_width": 1.0},
        **kwargs
    ):
        """
        Animate streamlines with flowing effects.
        
        Parameters:
        - stream_lines: StreamLines object to animate
        - lag_range: Stagger animation start times
        - rate_multiple: Animation speed multiplier
        - line_anim_config: Animation configuration for flow effects
        """

def move_along_vector_field(
    mobject: Mobject, 
    func: Callable[[Vect3], Vect3]
) -> Mobject:
    """
    Add updater to move mobject along vector field flow.
    
    Parameters:
    - mobject: Object to move
    - func: Vector field function
    
    Returns:
    Modified mobject with movement updater
    """

def move_points_along_vector_field(
    mobject: Mobject, 
    func: Callable, 
    coordinate_system: CoordinateSystem
) -> Mobject:
    """
    Move all points of a mobject along vector field.
    
    Parameters:
    - mobject: Object whose points will move
    - func: Vector field function
    - coordinate_system: Coordinate system for mapping
    
    Returns:
    Modified mobject with point movement updater
    """

def get_sample_coords(
    coordinate_system: CoordinateSystem, 
    density: float = 1.0
) -> np.ndarray:
    """
    Generate sample coordinates for vector field evaluation.
    
    Parameters:
    - coordinate_system: Coordinate system to sample
    - density: Sampling density
    
    Returns:
    Array of coordinate points for field evaluation
    """

def vectorize(pointwise_function: Callable) -> Callable:
    """
    Convert pointwise function to vectorized form for efficient evaluation.
    
    Parameters:
    - pointwise_function: Function operating on single points
    
    Returns:
    Vectorized function operating on arrays
    """

def ode_solution_points(
    function, 
    state0, 
    time, 
    dt: float = 0.01
) -> np.ndarray:
    """
    Solve ODE system to generate solution points.
    
    Parameters:
    - function: Derivative function for ODE system
    - state0: Initial state
    - time: Time duration for solution
    - dt: Time step
    
    Returns:
    Array of solution points
    """

from manimlib import *

class ElectricField(Scene):
    def construct(self):
        plane = NumberPlane(x_range=[-4, 4], y_range=[-3, 3])
        self.add(plane)
        
        # Define electric field from point charges
        def electric_field(coords):
            x, y = coords.T
            
            # Positive charge at (-1, 0)
            r1_squared = (x + 1)**2 + y**2 + 0.1  # Small offset to avoid singularity
            E1_x = (x + 1) / r1_squared**(3/2)
            E1_y = y / r1_squared**(3/2)
            
            # Negative charge at (1, 0)
            r2_squared = (x - 1)**2 + y**2 + 0.1
            E2_x = -(x - 1) / r2_squared**(3/2)
            E2_y = -y / r2_squared**(3/2)
            
            # Total field
            Ex = E1_x + E2_x
            Ey = E1_y + E2_y
            
            return np.array([Ex, Ey]).T
        
        # Create vector field
        field = VectorField(
            electric_field,
            plane,
            density=1.5,
            stroke_width=2,
            max_vect_len=0.6
        )
        
        # Add charges
        positive_charge = Circle(radius=0.1, color=RED, fill_opacity=1)
        negative_charge = Circle(radius=0.1, color=BLUE, fill_opacity=1)
        positive_charge.move_to(plane.c2p(-1, 0))
        negative_charge.move_to(plane.c2p(1, 0))
        
        # Labels
        plus_label = Text("+", font_size=24, color=WHITE).move_to(positive_charge)
        minus_label = Text("−", font_size=24, color=WHITE).move_to(negative_charge)
        
        self.play(ShowCreation(field), run_time=3)
        self.play(
            ShowCreation(positive_charge),
            ShowCreation(negative_charge),
            Write(plus_label),
            Write(minus_label)
        )
        
        # Add field lines
        field_lines = StreamLines(
            electric_field,
            plane,
            density=0.8,
            stroke_width=1.5,
            color_by_magnitude=True
        )
        
        self.play(ShowCreation(field_lines), run_time=4)
        self.wait(2)

class FluidFlow(Scene):
    def construct(self):
        plane = NumberPlane(x_range=[-3, 3], y_range=[-2, 2])
        self.add(plane)
        
        # Define fluid velocity field (vortex + uniform flow)
        def velocity_field(coords):
            x, y = coords.T
            
            # Vortex component
            vortex_x = -y
            vortex_y = x
            
            # Uniform flow component
            uniform_x = np.ones_like(x) * 0.5
            uniform_y = np.zeros_like(y)
            
            # Combine components
            vx = vortex_x + uniform_x
            vy = vortex_y + uniform_y
            
            return np.array([vx, vy]).T
        
        # Create streamlines
        streamlines = StreamLines(
            velocity_field,
            plane,
            density=1.5,
            n_repeats=2,
            noise_factor=0.1,
            stroke_width=2,
            color_by_magnitude=True,
            magnitude_range=(0, 3)
        )
        
        # Animate the flow
        animated_lines = AnimatedStreamLines(
            streamlines,
            lag_range=2,
            rate_multiple=1.5
        )
        
        self.add(animated_lines)
        
        # Add some particles that follow the flow
        particles = VGroup(*[
            Dot(radius=0.05, color=YELLOW).move_to(
                plane.c2p(
                    3 * (np.random.random() - 0.5),
                    2 * (np.random.random() - 0.5)
                )
            )
            for _ in range(20)
        ])
        
        # Make particles follow the field
        for particle in particles:
            move_along_vector_field(particle, lambda p: velocity_field(np.array([plane.p2c(p)[:2]]))[0])
        
        self.add(particles)
        self.wait(10)

class PhasePortrait(Scene):
    def construct(self):
        plane = NumberPlane(x_range=[-3, 3], y_range=[-3, 3])
        plane.add_coordinates()
        self.add(plane)
        
        # Simple harmonic oscillator phase space
        def harmonic_flow(coords):
            x, y = coords.T  # x = position, y = velocity
            dx_dt = y       # dx/dt = velocity
            dy_dt = -x      # dy/dt = -x (acceleration)
            
            return np.array([dx_dt, dy_dt]).T
        
        # Create vector field
        field = VectorField(
            harmonic_flow,
            plane,
            density=1.2,
            stroke_width=1.5,
            max_vect_len=0.4,
            color_map_name="viridis"
        )
        
        # Create phase trajectories
        trajectories = StreamLines(
            harmonic_flow,
            plane,
            density=0.6,
            stroke_width=2,
            solution_time=2*PI,  # One period
            color_by_magnitude=False,
            stroke_color=YELLOW
        )
        
        self.play(ShowCreation(field), run_time=2)
        self.play(ShowCreation(trajectories), run_time=3)
        
        # Add labels
        x_label = Text("Position", font_size=24).next_to(plane.x_axis, DOWN)
        y_label = Text("Velocity", font_size=24).next_to(plane.y_axis, LEFT)
        y_label.rotate(PI/2)
        
        title = Text("Harmonic Oscillator Phase Portrait", font_size=32).to_edge(UP)
        
        self.play(Write(title), Write(x_label), Write(y_label))
        self.wait()

class TimeVaryingField(Scene):
    def construct(self):
        plane = NumberPlane(x_range=[-2, 2], y_range=[-2, 2])
        self.add(plane)
        
        # Time-varying magnetic field (rotating)
        def rotating_field(coords, t):
            x, y = coords.T
            
            # Rotating uniform field
            field_strength = 1.0
            omega = 1.0  # Angular frequency
            
            Bx = field_strength * np.cos(omega * t) * np.ones_like(x)
            By = field_strength * np.sin(omega * t) * np.ones_like(y)
            
            return np.array([Bx, By]).T
        
        # Create time-varying field
        field = TimeVaryingVectorField(
            rotating_field,
            plane,
            density=1.5,
            stroke_width=3,
            max_vect_len=0.8
        )
        
        # Add field indicator
        field_arrow = Arrow(ORIGIN, RIGHT, color=RED, buff=0)
        field_arrow.to_edge(UP + RIGHT)
        
        def update_indicator(arrow):
            t = field.time
            direction = np.array([np.cos(t), np.sin(t), 0])
            arrow.become(Arrow(ORIGIN, direction, color=RED, buff=0))
            arrow.to_edge(UP + RIGHT)
        
        field_arrow.add_updater(update_indicator)
        
        # Labels
        title = Text("Rotating Magnetic Field", font_size=32).to_edge(UP + LEFT)
        time_label = Text("t = 0.0", font_size=24).to_edge(DOWN + RIGHT)
        
        def update_time_label(label):
            t = field.time
            label.become(Text(f"t = {t:.1f}", font_size=24))
            label.to_edge(DOWN + RIGHT)
        
        time_label.add_updater(update_time_label)
        
        self.add(field, field_arrow, title, time_label)
        self.wait(8)  # Let field rotate

from manimlib.mobject.interactive import LinearNumberSlider

class InteractiveField(Scene):
    def setup(self):
        # Create parameter controls
        self.strength_slider = LinearNumberSlider(
            value=1.0, min_value=0.1, max_value=3.0, step=0.1
        )
        self.rotation_slider = LinearNumberSlider(
            value=0.0, min_value=-PI, max_value=PI, step=0.1
        )
        
        self.strength_slider.to_edge(DOWN).shift(UP * 0.5)
        self.rotation_slider.to_edge(DOWN)
        
        self.add(self.strength_slider, self.rotation_slider)
    
    def construct(self):
        plane = NumberPlane(x_range=[-2, 2], y_range=[-2, 2])
        
        # Interactive field function
        def interactive_field(coords):
            x, y = coords.T
            strength = self.strength_slider.get_value()
            angle = self.rotation_slider.get_value()
            
            # Rotated uniform field
            cos_a, sin_a = np.cos(angle), np.sin(angle)
            field_x = strength * cos_a * np.ones_like(x)
            field_y = strength * sin_a * np.ones_like(y)
            
            return np.array([field_x, field_y]).T
        
        # Create responsive field
        field = VectorField(
            interactive_field,
            plane,
            density=1.5,
            stroke_width=2
        )
        
        # Add updater to redraw field when parameters change
        field.add_updater(lambda f: f.update_vectors())
        
        # Labels
        strength_label = Text("Field Strength", font_size=20)
        rotation_label = Text("Field Rotation", font_size=20)
        
        strength_label.next_to(self.strength_slider, LEFT)
        rotation_label.next_to(self.rotation_slider, LEFT)
        
        self.add(plane, field, strength_label, rotation_label)
        self.wait(15)  # Interactive exploration time

# Vectorized field functions for efficiency
def optimized_field(coords):
    # Use numpy operations on entire arrays
    x, y = coords.T
    return np.stack([np.sin(x) * np.cos(y), np.cos(x) * np.sin(y)], axis=1)

# Custom color mapping for large fields
def custom_color_map(magnitudes):
    # Efficient color mapping using numpy
    normalized = magnitudes / np.max(magnitudes)
    return plt.cm.plasma(normalized)

# Works with any coordinate system
polar_plane = PolarPlane()
field_polar = VectorField(radial_field, polar_plane)

# 3D coordinate systems
axes_3d = ThreeDAxes()
# Note: Vector fields are primarily 2D, but can work with projections