티스토리 수익 글 보기 키워드마스터

티스토리 수익 글 보기

gallery-examples/notebooks/math/ode-solver.py at main · marimo-team/gallery-examples · GitHub

History

164 lines (128 loc) · 3.92 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

# /// script

# requires-python = ">=3.12"

# dependencies = [

# "marimo",

# "matplotlib",

# "numpy",

# "scipy",

# "wigglystuff",

# ]

# ///

import marimo

__generated_with = "0.20.4"

app = marimo.App(width="medium")

@app.cell

def _():

import marimo as mo

return (mo,)

@app.cell

def _():

from wigglystuff.chart_puck import ChartPuck

import matplotlib.pyplot as plt

import numpy as np

from scipy.integrate import solve_ivp

return ChartPuck, np, plt, solve_ivp

@app.cell(hide_code=True)

def _(mo):

mo.md(r"""

# Interactive ODE Solver

Drag the red initial-condition point to explore solutions of

$$\frac{dy}{dt} = \frac{1}{2}y^2 – 5t^2 + 1$$

The solver integrates forward (blue) and backward (orange) in time

from the current initial condition.

""")

return

@app.cell

def _(ChartPuck, np, ode_rhs, plt, solve_ivp):

X_MIN, X_MAX = –3, 3

Y_MIN, Y_MAX = –10, 10

def draw_ode_soln(ax: plt.Axes, widget: ChartPuck) -> None:

x, y = widget.x[0], widget.y[0]

# — Direction field (quiver plot) —

_t_grid: np.ndarray = np.linspace(X_MIN, X_MAX, 25)

_y_grid: np.ndarray = np.linspace(Y_MIN, Y_MAX, 25)

_T, _Y = np.meshgrid(_t_grid, _y_grid) # Shape (25, 25) each

_dT: np.ndarray = np.ones_like(_T) # dt/dt = 1

_dY: np.ndarray = ode_rhs(_T, _Y).squeeze() # dy/dt from the ODE

# Normalize arrows so they all have the same length

_speed: np.ndarray = np.sqrt(_dT**2 + _dY**2)

_dT_norm: np.ndarray = _dT / _speed

_dY_norm: np.ndarray = _dY / _speed

ax.quiver(

_T,

_Y,

_dT_norm,

_dY_norm,

_speed,

cmap="coolwarm",

alpha=0.45,

scale=30,

width=0.003,

headwidth=3.5,

headlength=4,

)

# — Compute forward ODE solution —

t_forward_span: tuple[float, float] = (x, float(X_MAX))

sol_forward = solve_ivp(

ode_rhs, t_forward_span, [y], dense_output=True, max_step=0.05

)

t_fwd: np.ndarray = np.linspace(x, float(X_MAX), 300)

y_fwd: np.ndarray = sol_forward.sol(t_fwd)[0]

# — Compute backward ODE solution —

t_backward_span: tuple[float, float] = (x, float(X_MIN))

sol_backward = solve_ivp(

ode_rhs, t_backward_span, [y], dense_output=True, max_step=0.05

)

t_bwd: np.ndarray = np.linspace(x, float(X_MIN), 300)

y_bwd: np.ndarray = sol_backward.sol(t_bwd)[0]

# — Plot forward and backward solutions —

ax.plot(

t_fwd, y_fwd, color="tab:blue", linewidth=2.5, label="Forward", zorder=3

)

ax.plot(

t_bwd,

y_bwd,

color="tab:orange",

linewidth=2.5,

label="Backward",

zorder=3,

)

ax.set_title(f"Position: ({x:.2f}, {y:.2f})")

ax.set_xlim(X_MIN, X_MAX)

ax.set_ylim(Y_MIN, Y_MAX)

ax.set_xlabel("t")

ax.set_ylabel("y")

ax.legend(loc="upper left", fontsize=8)

ax.grid(True, alpha=0.3)

ax.figure.tight_layout()

return (draw_ode_soln,)

@app.cell

def _(np):

def ode_rhs(t: float, y: np.ndarray) -> np.ndarray:

"""dy/dt = 0.5 y^2 – 5t^2 + 1."""

return np.array([0.5 * y**2 – 5 * t**2 + 1])

return (ode_rhs,)

@app.cell

def _(ChartPuck, draw_ode_soln):

dynamic_puck = ChartPuck.from_callback(

draw_fn=draw_ode_soln,

x_bounds=(–3, 3),

y_bounds=(–5, 5),

figsize=(10, 6),

x=0,

y=0,

puck_color="red",

puck_radius=5,

throttle=10,

)

return (dynamic_puck,)

@app.cell

def _(dynamic_puck, mo):

dynamic_widget = mo.ui.anywidget(dynamic_puck)

return (dynamic_widget,)

@app.cell

def _(dynamic_widget):

dynamic_widget

return

if __name__ == "__main__":

app.run()

티스토리 수익 글 보기