second formulation

#!/usr/bin/env python3

###############################################################################
# Optimised MP4 renderer for the newer jellyfish-like formula
#
# Newer formula:
#
# a=(m,d=mag(k=9*cos(i/61),e=i/652-13)**2/89+1)=>point(
#   (q=79-e/2*sin(k)+k/d*(6+5*sin(sin(d*d+e/9-t+m))))*
#     sin(c=d/1.9+cos(t-d*3+m)/11-t/16+m)+200,
#   (q+40)*cos(c)+200
# )
#
# t=0,draw=$=>{
#   t||createCanvas(w=400,w);
#   background(9).stroke(w,96);
#   for(t+=PI/45,i=2e4;i--;)a(i%2*3)
# }
#
# This version:
#
#   - uses the newer formula;
#   - updates one existing scatter object frame by frame;
#   - writes directly to MP4 using ffmpeg;
#   - avoids writing thousands of intermediate PNG files;
#   - precomputes immutable quantities outside compute_frame().
###############################################################################

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, FFMpegWriter

# =============================================================================
# 1. Tunable parameters
# =============================================================================

n_points = 20_000

# 540 * 4 frames at 25 fps gives:
#
#   2160 / 25 = 86.4 seconds
n_frames = 540

fps = 25

# Mathematical canvas from the original p5.js code:
#
#   createCanvas(w = 400, w)
canvas_width = 400
canvas_height = 400

# Rendered output size.
render_width_px = 600
render_height_px = 600

# The final pixel size is figsize_inches * dpi.
dpi = 100

# Visual tuning.
#
# Matplotlib scatter "s" is marker area in points squared.
point_size = 1.5
point_alpha = 0.70

background_colour = "#090909"

output_file = "jellyfish_new_formula_python.mp4"

# FFmpeg settings.
ffmpeg_crf = 18
ffmpeg_preset = "veryfast"
ffmpeg_threads = 4

# =============================================================================
# 2. Static point data
# =============================================================================

# Point index:
#
#   i = 0, 1, ..., n_points - 1
#
# These are not x-coordinates. They are indices used to generate the point cloud.
i = np.arange(n_points)

# New formula:
#
#   k = 9 * cos(i / 61)
k = 9 * np.cos(i / 61)

# New formula:
#
#   e = i / 652 - 13
e = i / 652 - 13

# New formula:
#
#   d = mag(k, e)^2 / 89 + 1
#
# Since:
#
#   mag(k, e)^2 = k^2 + e^2
#
# we compute it directly as:
d = (k**2 + e**2) / 89 + 1

# New formula:
#
#   m = i % 2 * 3
#
# This alternates between 0 and 3 and creates the two intertwined creatures.
m = (i % 2) * 3

# =============================================================================
# 3. Precomputed immutable terms
# =============================================================================

# These do not depend on frame/time, so compute them once.

d2 = d**2

# Static part of q:
#
#   79 - e/2*sin(k)
q_static = 79 - e / 2 * np.sin(k)

# Multiplier in q:
#
#   k / d
q_multiplier = k / d

# Static part inside the nested sine:
#
#   d*d + e/9 - t + m
#
# Frame-dependent part is just "- t".
wave_static = d2 + e / 9 + m

# Static part of c:
#
#   d/1.9 + cos(t - d*3 + m)/11 - t/16 + m
#
# Split into:
#
#   c_static = d/1.9 + m
#   cosine argument = t + (-d*3 + m)
c_static = d / 1.9 + m
c_wave_static = -d * 3 + m

# Preallocated coordinate array for Matplotlib scatter offsets.
#
# Matplotlib expects an N x 2 array:
#
#   [[x1, y1],
#    [x2, y2],
#    ...]
offsets = np.empty((n_points, 2), dtype=np.float64)

# =============================================================================
# 4. Frame computation
# =============================================================================

def compute_frame(frame_id):
    """
    Compute one animation frame.

    Parameters
    ----------
    frame_id : int
        Integer frame number: 1, 2, 3, ...

        It is converted to continuous time as:

            t = frame_id * pi / 45

        matching the original JavaScript update:

            t += PI / 45

    Returns
    -------
    offsets : numpy.ndarray
        N x 2 array with x/y coordinates for all points.
    """

    # Frame time.
    t = frame_id * np.pi / 45

    # -------------------------------------------------------------------------
    # Main deformation term q
    # -------------------------------------------------------------------------
    #
    # Original:
    #
    #   q = 79 - e/2*sin(k) +
    #       k/d * (6 + 5*sin(sin(d*d + e/9 - t + m)))
    #
    q = (
        q_static
        + q_multiplier
        * (6 + 5 * np.sin(np.sin(wave_static - t)))
    )

    # -------------------------------------------------------------------------
    # Angular term c
    # -------------------------------------------------------------------------
    #
    # Original:
    #
    #   c = d/1.9 + cos(t - d*3 + m)/11 - t/16 + m
    #
    c = (
        c_static
        + np.cos(t + c_wave_static) / 11
        - t / 16
    )

    # -------------------------------------------------------------------------
    # Final coordinates
    # -------------------------------------------------------------------------
    #
    # Original:
    #
    #   x = q * sin(c) + 200
    #   y = (q + 40) * cos(c) + 200
    #
    offsets[:, 0] = q * np.sin(c) + 200
    offsets[:, 1] = (q + 40) * np.cos(c) + 200

    return offsets

# =============================================================================
# 5. Matplotlib figure setup
# =============================================================================

fig_width_inches = render_width_px / dpi
fig_height_inches = render_height_px / dpi

fig, ax = plt.subplots(
    figsize=(fig_width_inches, fig_height_inches),
    dpi=dpi
)

fig.patch.set_facecolor(background_colour)
ax.set_facecolor(background_colour)

ax.set_xlim(0, canvas_width)
ax.set_ylim(0, canvas_height)
ax.set_aspect("equal")

ax.set_xticks([])
ax.set_yticks([])

for spine in ax.spines.values():
    spine.set_visible(False)

ax.margins(0)
plt.subplots_adjust(left=0, right=1, top=1, bottom=0)

# =============================================================================
# 6. Initial frame and scatter object
# =============================================================================

initial_offsets = compute_frame(1).copy()

points = ax.scatter(
    initial_offsets[:, 0],
    initial_offsets[:, 1],
    s=point_size,
    c="white",
    alpha=point_alpha,
    linewidths=0
)

# =============================================================================
# 7. Animation update function
# =============================================================================

def update(frame_id):
    current_offsets = compute_frame(frame_id)

    # Update the existing scatter object.
    # This avoids rebuilding the plot for every frame.
    points.set_offsets(current_offsets)

    return (points,)

# =============================================================================
# 8. Build and save animation directly as MP4
# =============================================================================

animation = FuncAnimation(
    fig,
    update,
    frames=np.arange(1, n_frames + 1),
    interval=1000 / fps,
    blit=True
)

writer = FFMpegWriter(
    fps=fps,
    codec="libx264",
    bitrate=-1,
    extra_args=[
        "-preset", ffmpeg_preset,
        "-crf", str(ffmpeg_crf),
        "-threads", str(ffmpeg_threads),
        "-pix_fmt", "yuv420p",
    ],
)

animation.save(
    output_file,
    writer=writer
)

plt.close(fig)

print(f"Saved: {output_file}")