```python
%matplotlib inline
%run ../setup/nb_setup
```

# Orbits 1: Orbits in Spherical Gravitational Potentials

Author(s): Adrian Price-Whelan, Sarah Pearson

## Learning goals

The main goal for this tutorial is to introduce the concept of orbits and numerically-computed trajectories using spherical gravitational potential models.

## Introduction

For many contexts in Galactic dynamics, the orbital trajectory of an object (star, planet, star cluster, satellite galaxy, etc.) is a fundamental tool for studying and interpreting the dynamical state, evolution, or initial conditions of a gravitational system. However, in general, we cannot directly observe any appreciable changes in position or velocity of a gravitating body over human timescales. As Galactic dynamicists, we therefore effectively only have access to measurements of a single kinematic snapshot of any dynamical system. With a kinematic snapshot and a model for the overall mass distribution of a system, we then seek to compute the past or future trajectories of the constituent bodies and use this information to, for example, refine our mass model, interpret kinematic features, or predict the evolution (past or present) of a system. In many such regimes, the mass of the orbiting body is small relative to the larger mass distribution in which it orbits (e.g., a star within a Galaxy). It is therefore often sufficient to treat an orbiting body as an effectively massless "tracer" particle whose trajectory is set by the broader mass distribution, and which does not cause any appreciable backreaction of the host system. (In later workshops, we will explore cases in which this assumption breaks down, such as with the Milky Way–LMC interaction.) In this context, a common goal is to compute the time-dependent position and velocity — the orbit — of an effectively massless body orbiting within an external gravitational field.

While galaxies are actually composed of a large number of discrete point masses and structures (and we do not know the smallest size of dark matter substructures), it is often sufficient to model the mass distribution of a galaxy by considering smooth mass distributions (where the smoothing length scale is much smaller than the size of a galaxy, <<100 kpc, but larger than the spacing between massive structures, >>1 pc). As we saw in the lectures, and as you might recall from past Mechanics courses, the fundamental quantity in the context of computing orbits is the acceleration field $\boldsymbol{a}(\boldsymbol{x})$ specified often as a function of a three-dimensional position $\boldsymbol{x}$. In convervative systems — or effectively conservative systems: neglecting star formation and death/supernovae for stars, and assuming that dark matter acts effectively like a collisionless fluid, the massive objects that make up a galaxy ineract only through gravitational interactions — the (vector) acceleration field can be directly computed as the gradient of the (scalar) gravitational potential $\Phi(\boldsymbol{x})$ of a system:

$$
\boldsymbol{a}(\boldsymbol{x}) = -\nabla \Phi(\boldsymbol{x})
$$

where the gravitational potential (sometimes later abbreviated simply as "potential") is related to the mass (density) distribution through Poisson's equation:

$$
\nabla^2 \Phi(\boldsymbol{x}) = 4\pi \, G \, \rho(\boldsymbol{x})
$$

While the acceleration or gravitational field is "closer" to the dynamics (in that it is the force on a body that determines how it responds) and the density field is a direct representation of the mass distribution, it is often easier to work with and visualize the gravitational potential.

In this tutorial, we will introduce some commonly-used spherical gravitational potential models, compute the orbits of particles in some of these models, and analyze the properties of orbits in spherical models.

## Terminology and Notation

- 3D position: $\boldsymbol{x}$
- 3D velocity: $\boldsymbol{v}$
- 6D phase-space position: $\boldsymbol{w} = (\boldsymbol{x}, \boldsymbol{v})$
- Orbit: $\boldsymbol{w}(t) = (\boldsymbol{x}(t), \boldsymbol{v}(t))$
- Gravitational potential: $\Phi(\boldsymbol{x})$
- Mass density: $\rho(\boldsymbol{x})$
- Spherical radius: $r = |\boldsymbol{x}|$
- Unit vector: $\hat{\boldsymbol{x}}$
- Time derivative: $\dot{x} = \textrm{d}x/\textrm{d}t$
- TODO: review whether I should put other terminology here!

### Notebook Setup and Package Imports

```python
from astropy.constants import G
import astropy.coordinates as coord
import astropy.units as u
from astropy.time import Time

from IPython.display import display, Math
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import solve_ivp
from scipy.signal import argrelmax
import sympy as sy
```

## Orbits in a Point Mass Potential

One of the simplest spherical mass models is that of a point mass, in which:

$$
\Phi(r) = - \frac{G \, M}{r}
$$

or equivalently

$$
\Phi(\boldsymbol{x}) = - \frac{G \, M}{|\boldsymbol{x}|}
$$

and the acceleration field is

$$
\begin{align}
\boldsymbol{a}(\boldsymbol{x}) &= -\nabla \Phi(\boldsymbol{x}) \\
    & = - \frac{G \, M}{|\boldsymbol{x}|^3} \, \boldsymbol{x} \\
    & = - \frac{G \, M}{r^2} \, \hat{\boldsymbol{x}}
\end{align}
$$

Let's implement a function that computes the point mass acceleration at any 3D position, given a mass M:

```python
def ptmass_acceleration(xyz, M):
    """
    This function computes the 3D acceleration vector at a
    position or an array of positions "xyz" from a point mass
    with mass "M" at the origin.

    Parameters
    ----------
    xyz : quantity-like
        This should have shape (3, ...)
    M : quantity-like
        The total mass of the point mass.

    """
    r3 = np.sum(xyz ** 2, axis=0) ** (3 / 2)
    dPhi_dx = G * M * xyz / r3
    return -dPhi_dx
```

In the function definition above, note that we use the gravitational constant `G` as imported from `astropy.constants`. This object is more than just a number: It is a special kind of Python/Astropy object called a [`Quantity`](https://docs.astropy.org/en/stable/units/quantity.html#quantity), which carries along and propagates units. For example, if we view the representation of the gravitational constant:

```python
G
```

we see that it also displays the units $\frac{\textrm{m}^3}{\textrm{kg} \, \textrm{s}^2}$! When we call the `ptmass_acceleration()` function, this means that we should pass in a position and a point mass with units, and we will get out an acceleration that we can easily convert to other unit representations. For example, assuming the Earth is on the x-axis in some solar system coordinate system, we can evaluate the acceleration felt on Earth by the Sun as:

```python
earth_pos = [1, 0, 0] * u.au  # 1 Astronomical Unit along the x axis
earth_acc = ptmass_acceleration(earth_pos, M=1 * u.Msun)
earth_acc
```

This is a strange unit! By default, `astropy.units` propagates units but does not do any conversion or simplification of the units. For example, here the $\textrm{M}_\odot$ in the numerator should cancel the kg in the denominator. If we want to convert this into something more readable, we use the `.to()` method and specify the desired output unit:

```python
earth_acc.to(u.km / u.s ** 2)
```

Some other examples of acceleration units:

```python
earth_acc.to(u.au / u.yr ** 2)
```

```python
earth_acc.to(u.km / u.s / u.day)
```

### Example: Compute the acceleration from the Sun at the current (right now!) location of Mars

Get a representation of the current time using `astropy.time.Time()`:

```python
now = Time.now()
now.mjd
```

We will use Astropy (`astropy.coordinates`) to retrieve the position of Mars: In the background, this queries the JPL ephemeris for us. So we can retrieve the position of Mars at the current time using:

```python
mars_pos = coord.get_body("mars", now)
mars_pos
```

(this works for other solar system bodies - try it! Swap 'mars' out for, e.g., 'saturn')

By default, the position of Mars is returned to us in a Geocentric (Earth-centric) coordinate system. We need to convert this to a Heliocentric coordinate system so we can compute the acceleration from the Sun (assuming it is at the origin). We will convert it to Heliocentric Ecliptic coordinates:

```python
mars_ecl = mars_pos.transform_to(coord.HeliocentricTrueEcliptic(obstime=now))
```

Now we can retrieve a Cartesian representation of the Ecliptic position of Mars, and use this to compute the acceleration:

```python
mars_xyz_now = mars_ecl.cartesian.xyz
mars_xyz_now
```

```python
mars_acc = ptmass_acceleration(mars_xyz_now, M=1 * u.Msun)
mars_acc.to(u.km / u.s / u.day)
```

### Exercise: Compute the acceleration on the Moon from the Sun and Earth - which is larger?

*Hint: The Astropy unit `u.Mearth` might be useful!*

```python

```

---

### Computing an orbit in the point mass potential

The orbit of a massless particle in any acceleration field can be computed by solving the differential equation

$$
\begin{align}
\ddot{\boldsymbol{x}} &= \boldsymbol{a}(\boldsymbol{x}) \\
\ddot{\boldsymbol{x}} &= -\nabla \Phi(\boldsymbol{x})
\end{align}
$$

For a point mass, this is sometimes referred to as the "Kepler problem": This is a special case in which we could write down an analytic expression for the position of the particle at any time, given a set of initial conditions. However, generically, this differential equation cannot be solved analytically and we must instead resort to numerical integration tools. In this context, it is often easier to reduce the 2nd-order differential equation for postion expressed above into a pair of coupled 1st-order differential equations for position and velocity (or momentum). We do this be introducing an auxiliary variable to represent the 3D velocity, so that the relevant system of differential equations we need to solve is now:

$$
\begin{align}
\dot{\boldsymbol{x}} &= \boldsymbol{v} \\
\dot{\boldsymbol{v}} &= -\nabla \Phi(\boldsymbol{x})
\end{align}
$$

We can compress this into a single expression by defining the "phase-space position" $\boldsymbol{w} = (\boldsymbol{x}, \boldsymbol{v})$ such that:

$$
\begin{align}
\dot{\boldsymbol{w}} &= (\dot{\boldsymbol{x}}, \dot{\boldsymbol{v}}) \\
&= (\boldsymbol{v}, -\nabla \Phi(\boldsymbol{x}))
\end{align}
$$

Our goal then is to compute $\boldsymbol{w}(t)$ given initial conditions $\boldsymbol{w}_0$. There are many standard tools available for solving problems like this numerically. This is sometimes called an "initial value problem," or "ordinary differential equation" (ODE) integration. Here, we will define a function to compute the time-derivative of $\boldsymbol{w}$ and use `scipy.integrate.solve_ivp()` to compute the orbital trajectory of a particle ([solve_ivp() documentation](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html) for context).

The function we defined above to compute the point mass acceleration only does half of what we need: We need to define a new function `ptmass_F()` that takes the full phase-space position $\boldsymbol{w}$ and returns its time-derivative, $\dot{\boldsymbol{w}} = (\boldsymbol{v}, -\nabla \Phi(\boldsymbol{x}))$. `solve_ivp()` also requires that the function accept a first argument, the time `t`, which in this case we will include in the definition but ignore in any expressions.

Also, `solve_ivp()` unfortunately does not support `astropy.units`, so we will have to re-implement the point mass acceleration in a way that takes only regular Python numbers. To avoid having to specify units for both `G` and the mass `M`, we will have this function accept the combined parameter `GM = G * M`:

```python
def ptmass_F(t, w, GM):
    """
    Parameters
    ----------
    t : numeric
        The current time during the integration.
    w : array-like
        The phase-space position, which should have 6 elements.
    GM : numeric
        Mass times the gravitational constant.

    """
    xyz = w[:3]
    vxyz = w[3:]

    r3 = np.sum(xyz ** 2, axis=0) ** (3 / 2)
    dPhi_dx = GM * xyz / r3
    return np.concatenate([vxyz, -dPhi_dx])
```

Even though we cannot use `astropy.units` to keep track of units internally, we can still use it to do unit conversions before we pass inputs to the function, or with the resulting orbital time series that we compute. For this, it helps to work in a consistent unit system (a set of base units that represent length, mass, and time). Here, since we are using Solar System examples, a good unit system is AU, $\textrm{M}_\odot$, and years:

```python
unitsys = [u.au, u.Msun, u.year]
```

### Circular orbits and the circular velocity curve

We now have all of the pieces we need to compute an orbit in the point mass potential! How then should we choose initial conditions for a trial orbit? One simple option is to pick a location, and start a particle on a circular orbit that goes through that location.

An important property of spherical mass distributions is the concept of the *circular velocity* $v_{\rm circ}(r)$, which is defined to be the speed of a test particle on a circular orbit at radius $r$. The circular velocity can be found by equating the centripetal acceleration $\frac{v^2}{r}$ with the (absolute value of the) acceleration from the gravitational potential $\nabla \Phi$

$$
\begin{align}
\frac{v^2}{r} &= \left|\frac{\textrm{d}\Phi}{\textrm{d}r}\right| \\
v_{\rm circ}(r) &= \sqrt{r \, \left|\frac{\textrm{d}\Phi}{\textrm{d}r}\right|}
\end{align}
$$

For the point mass potential, this is simply:

$$
v_{\rm circ}(r) = \sqrt{\frac{G\,M}{r}}
$$

### Example: Compute the circular velocity in at $r=1.6~\textrm{AU}$ in the Solar System in AU/year

```python
r0 = 1.6 * u.au
M = 1 * u.Msun
v_circ = np.sqrt(G * M / r0).decompose(unitsys)
v_circ
```

The `.decompose()` method allows us to represent a `Quantity` in a given unit system without having to construct the unit explicitly (here we just had to pass in the list of base units to `decompose()`).

We will now use this radius and circular velocity to construct 3D initial conditions that we will pass to `solve_ivp()` to compute the orbit of a test particle in the Solar System. We will start the particle on the x-axis ($y=z=0$) so that the velocity (perpendicular to the position vector) is purely in the $\hat{y}$ direction:

```python
x0 = [1, 0, 0] * r0
v0 = [0, 1, 0] * v_circ
w0 = np.concatenate((x0.decompose(unitsys).value, v0.decompose(unitsys).value))
w0
```

Finally, we have to convert the combined amplitude parameter `GM` into our unit system:

```python
GM = (G * M).decompose(unitsys).value
GM
```

The last thing we need to specify is the grid of times over which to compute the orbit. For this example, we will use 1024 equally spaced points between 0 and 10 years, the time unit in our unit system.

```python
t_eval = np.linspace(0, 10, 1024)
```

```python
soln = solve_ivp(
    ptmass_F, y0=w0, args=(GM,), t_span=(t_eval.min(), t_eval.max()), t_eval=t_eval
)
```

The returned Solution object `soln` now contains the orbit of our test particle (position and velocity) evaluated at all 1024 points in our time grid!

```python
soln
```

The Scipy integrators use the variable "`y`" for any integration variables, so in this case `w = y`. To plot the position time series of the particle, we will plot two 2D projections of the 3D position: x-y and x-z:

```python
fig, axes = plt.subplots(
    1, 2, figsize=(10, 5), sharex=True, sharey=True, constrained_layout=True
)

# x-y
axes[0].plot(soln.y[0], soln.y[1])
axes[0].set_xlabel(f"$x$ [{unitsys[0]:latex_inline}]")
axes[0].set_ylabel(f"$y$ [{unitsys[0]:latex_inline}]")

# x-z
axes[1].plot(soln.y[0], soln.y[2])
axes[1].set_xlabel(f"$x$ [{unitsys[0]:latex_inline}]")
axes[1].set_ylabel(f"$z$ [{unitsys[0]:latex_inline}]")

for ax in axes:
    ax.set_aspect("equal")

fig.suptitle("default integrator (RK45)", fontsize=24)
```

The orbit is confined to the x-y plane, as we might expect (we only gave the particle velocity in the y direction, starting from the x axis). However: **that doesn't look like a circular orbit**: The orbit in the x-y plane looks more like a spiral!?! This is a telltale sign of energy drift. The force on our test particle is conservative and time-independent, so the orbital energy should be conserved along the particle's orbit. Let's compute the total energy for this numerically-integrated orbit and see if it is constant. To do that, we will implement a function to evaluate the point mass potential energy (defined above), and the total energy (the sum of potential and kinetic energy) at any phase-space position:

```python
def ptmass_potential(xyz, M):
    """
    Compute the potential energy from a point mass at the given
    position(s) ``xyz``.

    Parameters
    ----------
    xyz : quantity-like
    M : quantity-like
    """
    r = np.sqrt(np.sum(xyz ** 2, axis=0))
    return -G * M / r

def ptmass_total_energy(w, M, unitsys):
    """
    Compute the total energy (kinetic + potential) for phase-space
    position(s) ``w``.

    Parameters
    ----------
    w : array-like
    M : quantity-like
    unitsys : list, collection
        A list or iterable containing the base units that define the
        adopted unit system. Order: length, mass, time
    """
    xyz = w[:3] * unitsys[0]
    vxyz = w[3:] * unitsys[0] / unitsys[2]
    E_kin = 0.5 * np.sum(vxyz ** 2, axis=0)
    E_pot = ptmass_potential(xyz, M)
    return E_kin + E_pot
```

We can now compute the energy for each point in the orbital time series, and the fractional energy change relative to the initial energy `dE`:

```python
E = ptmass_total_energy(soln.y, M, unitsys)
dE = np.abs((E[1:] - E[0]) / E[0])

plt.figure(figsize=(6, 5))
plt.loglog(dE)
plt.xlabel("timestep number, $j$")
plt.ylabel("fractional energy drift\n" + r"$\left|(E_j - E_0) / E_0\right|$")
plt.ylim(1e-15, 1e0)
plt.suptitle("default integrator (RK45)", fontsize=24)
```

Aha! The energy of the computed orbit drifts steadily toward a nearly 100% change from the initial energy: This explains why the computed orbit does not look like a circular orbit, as we would expect. The fact that energy is not conserved is a consequence of the default integration method used by `solve_ivp()`: An (effectively) 4th-order Runge–Kutta method. Runge–Kutta methods are class of iterative schemes for computing approximate solutions for ordinary differential equations. However, they are approximate, and so these methods steadily accumulate errors. With `solve_ivp()`, we can control the error tolerance (the internal timestep used to perform the integration) using the function arguments `rtol` and `atol`, which set the threshold relative (rtol) and absolute (atol) error tolerances (see the [solve_ivp() documentation](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html) for more context).

Let's try re-running the orbit calculation, but with much smaller `rtol` and `atol` values (compared to the defaults of 1e-3 and 1e-6, respectively):

```python
soln_tol = solve_ivp(
    ptmass_F,
    y0=w0,
    args=(GM,),
    t_span=(t_eval.min(), t_eval.max()),
    t_eval=t_eval,
    rtol=1e-12,
    atol=1e-8,
)
```

```python
fig, axes = plt.subplots(
    1, 2, figsize=(10, 5), sharex=True, sharey=True, constrained_layout=True
)

# x-y
axes[0].plot(soln_tol.y[0], soln_tol.y[1])
axes[0].set_xlabel(f"$x$ [{unitsys[0]:latex_inline}]")
axes[0].set_ylabel(f"$y$ [{unitsys[0]:latex_inline}]")

# x-z
axes[1].plot(soln_tol.y[0], soln_tol.y[2])
axes[1].set_xlabel(f"$x$ [{unitsys[0]:latex_inline}]")
axes[1].set_ylabel(f"$z$ [{unitsys[0]:latex_inline}]")

for ax in axes:
    ax.set_aspect("equal")

fig.suptitle("RK45 (improved tolerance)", fontsize=24)
```

That looks much better! The orbit looks circular, and it visually appears to stay circular and planar for the duration of our integration window.

We therefore expect the energy conservation to be preserved much more accurately for this case:

```python
E = ptmass_total_energy(soln_tol.y, M, unitsys)
dE = np.abs((E[1:] - E[0]) / E[0])

plt.figure(figsize=(6, 5))
plt.loglog(dE)
plt.xlabel("timestep number, $j$")
plt.ylabel("fractional energy drift\n" + r"$\left|(E_j - E_0) / E_0\right|$")
plt.ylim(1e-15, 1e0)
plt.title("RK45 (improved tolerance)", fontsize=24)
```

Indeed, by the end of our integration window, the energy changes by $<10^{-7}$ relative to the initial energy for this case (the axis limits are the same here as in the previous plot! Scroll up to compare the two plots).

## Other Integration Algorithms

Runge–Kutta methods represent one class of numerical ODE integration tools, but there are several other classes of algorithms that are commonly used. Another class of algorithms that is used especially for evolving dynamical systems are *symplectic* integration methods. As we saw above with the Runge–Kutta scheme, over long times, errors introduced from approximations inherent to the method accumulate and lead to an increasing net change in the energy (which should be conserved). Symplectic integrators are special in that the iterative steps in the algorithm are *canonical transformations* themselves. A result of this is that, in conservative systems, over long times or averaged over time, these integrator classes conserve energy (in an average sense) and phase-space volume. In detail, the energy actually oscillates around the true energy, but it does not diverge from approximations in the method and only drifts due to numerical precision, which happens over much longer (integration) timescales.

One of the most well-known symplectic integration methods is the [*Leapfrog algorithm*](https://en.wikipedia.org/wiki/Leapfrog_integration), which is a specific type of more general Verlet integration schemes. Leapfrog is a 2nd-order method (the truncation error is 2nd-order) meaning that at fixed timestep, it will not be as precise as, for example, a 5th-order Runge–Kutta scheme. However, the approximate stability of the energy and the conserved phase volume is often desirable for long integrations. There are also 3rd, 4th, and higher-order analogs to Leapfrog integration that have better error properties (at fixed timestep) over short timescales.

### Leapfrog Integration

The default algorithm for Leapfrog integration advances the position and velocity (momentum) coordinates separately, and out of sync by 1/2 of a timestep, hence the "leapfrog" in the name: When the velocity advances by a full step from (timestep $i - \frac{1}{2}$) to (timestep $i + \frac{1}{2}$), it "leapfrogs" over the position (at timestep $i$). In a bit more detail, Leapfrog works as follows, starting from initial conditions $\boldsymbol{x}_0$ and $\boldsymbol{v}_0$ at time $t=t_0$:

- Compute the acceleration on the particle at the initial time:

  $$
  \boldsymbol{a}(t_0) = -\nabla \Phi|_{\boldsymbol{x}_0}
  $$

- Advance the velocity by 1/2 of a timestep $h=\Delta t$:

  $$
  \boldsymbol{v}_{1/2} = \boldsymbol{v}_0 + \frac{h}{2}\,\boldsymbol{a}(t_0)
  $$

- For each timestep $i$:

    - Increment the position by a full timestep using the "old" velocity, $\boldsymbol{v}_{i - 1/2}$

      $$
      \boldsymbol{x}_{i} = \boldsymbol{x}_{i-1} + h\,\boldsymbol{v}_{i - 1/2}
      $$

    - Compute the acceleration at the "new" position:

      $$
      \boldsymbol{a}(t_i) = -\nabla \Phi|_{\boldsymbol{x}_i}
      $$

    - Increment the velocity by a full timestep using the "new" acceleration

      $$
      \boldsymbol{v}_{i+1/2} = \boldsymbol{v}_{i-1/2} + h\,\boldsymbol{a}(t_i)
      $$

One caveat is that, if we stored all of the positions $\boldsymbol{x}_{i}$ and velocities $\boldsymbol{v}_{i+1/2}$, they would be out of sync with one another. In practice, it is common to instead split the velocity timestep increment into two increments by half of a timestep, so the last step in each iteration would look like:

- Increment the velocity by two half timesteps using the "new" acceleration
  $$
  \boldsymbol{v}_{i} = \boldsymbol{v}_{i-1/2} + \frac{h}{2}\,\boldsymbol{a}(t_i)
  \boldsymbol{v}_{i+1/2} = \boldsymbol{v}_{i} + \frac{h}{2}\,\boldsymbol{a}(t_i)
  $$

where we then save the (synchronized) $\boldsymbol{x}_{i}$ and $\boldsymbol{v}_{i}$ values as the orbital time series.

Let's implement this algorithm as a function, like `solve_ivp()`, which itself takes a function that computes the acceleration as the argument `acc_func`, along with initial conditions `x0` and `v0`, a timestep `dt`, and a number of timesteps `n_steps`. Because we are implementing this ourselves, we can build in support for Astropy units:

```python
def leapfrog_integrate(acc_func, x0, v0, dt, n_steps, acc_func_args=()):
    """
    Parameters
    ----------
    acc_func : callable
        A function that computes the acceleration at a position.
    x0 : quantity-like [length]
        Initial position.
    v0 : quantity-like [velocity]
        Initial velocity.
    dt : quantity-like [time]
        Timestep.
    n_steps : int
        The number of timesteps to take.
    acc_func_args : iterable, list (optional)
        Any additional arguments to be passed to the `acc_func` function.
    """
    # Make sure the input initial conditions and timestep are valid:
    x0 = u.Quantity(x0)
    v0 = u.Quantity(v0)
    ts = u.Quantity(dt) * np.arange(int(n_steps) + 1)

    # Initialize the velocity: step forward by 1/2 timestep
    a0 = acc_func(x0, *acc_func_args)
    v_1_2 = v0 + a0 * dt / 2.0

    # Set up array Quantity objects to store the full orbits for both
    # position and velocity
    xs = np.full((3, len(ts)), np.nan) * x0.unit
    xs[:, 0] = x0
    vs = np.full((3, len(ts)), np.nan) * v0.unit
    vs[:, 0] = v0

    x_im1 = x0
    v_im1_2 = v_1_2
    for i in range(n_steps):
        x_i = x_im1 + v_im1_2 * dt
        a_i = acc_func(x_i, *acc_func_args)

        v_i = v_im1_2 + a_i * dt / 2
        v_ip1_2 = v_i + a_i * dt / 2

        xs[:, i + 1] = x_i
        vs[:, i + 1] = v_i

        x_im1 = x_i
        v_im1_2 = v_ip1_2

    return ts, xs, vs
```

We can now compute the same orbit, using the same point mass potential and initial conditions as above. We will compute it for many more timesteps to compare the long-term behavior of Leapfrog with the Runge–Kutta schemes used above (this takes about 15 seconds on my machine - it is much slower because operations are slower with astropy Quantities as compared to numpy arrays - ask me about how we could speed this up if you are curious!):

```python
lf_ts, lf_xs, lf_vs = leapfrog_integrate(
    ptmass_acceleration, x0, v0, dt=1 * u.day, n_steps=16384, acc_func_args=(M,)
)
```

Let's now plot the same projections of the positional components of the orbit to compare with the plots above:

```python
fig, axes = plt.subplots(
    1, 2, figsize=(10, 5), sharex=True, sharey=True, constrained_layout=True
)

# x-y
axes[0].plot(lf_xs[0], lf_xs[1])
axes[0].set_xlabel(f"$x$ [{unitsys[0]:latex_inline}]")
axes[0].set_ylabel(f"$y$ [{unitsys[0]:latex_inline}]")

# x-z
axes[1].plot(lf_xs[0], lf_xs[2])
axes[1].set_xlabel(f"$x$ [{unitsys[0]:latex_inline}]")
axes[1].set_ylabel(f"$z$ [{unitsys[0]:latex_inline}]")

for ax in axes:
    ax.set_aspect("equal")

fig.suptitle("Leapfrog", fontsize=24)
```

This looks like the same circular orbit that we computed above! How well does it conserve energy?

```python
E = ptmass_potential(lf_xs, M) + 0.5 * np.sum(lf_vs ** 2, axis=0)
E = E.decompose(unitsys)
dE = np.abs((E[1:] - E[0]) / E[0])

plt.figure(figsize=(6, 5))
plt.loglog(dE)
plt.xlabel("timestep number, $j$")
plt.ylabel("fractional energy drift\n" + r"$\left|(E_j - E_0) / E_0\right|$")
plt.ylim(1e-15, 1e0)
plt.title("Leapfrog", fontsize=24)
```

Interestingly, as we expect, the error conservation is different: Rather than steadily rising toward larger numbers of timesteps (compare to the two equivalent figures above), the energy here is not constant, but appears to oscillate around a constant value. We can see this clearer by plotting the energy difference from the mean on a linear-linear plot:

```python
plt.figure(figsize=(6, 5))
mean_E = np.mean(E)
plt.plot((E - mean_E) / mean_E)
plt.xlabel("timestep number, $j$")
plt.ylabel("energy relative to mean\n" + r"$E_j - {\rm mean}_j(E_j)$")
plt.title("Leapfrog", fontsize=24)
plt.xlim(0, 10000)
```

The oscillations of the energy are small (on the order of $10^{-10}$ fractional changes in the energy), and appear more or less stable over the timescale that we integrated this orbit over.

We hope that this comparison of Runge–Kutta and Leapfrog methods for the point mass potential gives you a sense that there are often many tools for the job of numerical orbit computation. But which integration scheme should you use? In many contexts in Galactic dynamics, it doesn't make much of a difference: As long as your timestep is small enough that you conserve energy with enough precision for the problem at hand, you can use any of these integration schemes. However, for some problems it may be important to more precisely conserve phase-volume, in which case symplectic integrators are often a safer choice.

In practice, while it is fine to build your own machinery to do orbit integration, there are several tools available for Galactic dynamics that do a lot of the heavy lifting for you (and do it fast!). For example, in Python, your options are:

- [Gala](http://gala.adrian.pw/)
- [Galpy](https://github.com/jobovy/galpy)
- [Agama](http://agama.software/)

most of which implement the integration methods mentioned here (and others!), along with many of the gravitational potential models that we will discuss in the next few sections and tutorials.

### Exercise: Varying the initial conditions

So far, we have only worked with initial conditions for a circular orbit. Compute a few more orbits using `leapfrog_integrate()`, but with different choices for initial conditions (e.g., try different y velocity values, or try generating random velocity vectors). What do generic orbits look like?

```python

```

## The Plummer Potential

The Kepler potential is special because it represents a point mass, i.e. the density distribution is infinite at its origin and zero everywhere else. Most mass distributions that are relevant for Galactic dynamics are well-represented by density distributions that are finite everywhere or throughout some enclosed volume.

Mathematically, one of the simplest spherical extended (non-point) mass distributions is the [Plummer Sphere](https://en.wikipedia.org/wiki/Plummer_model), given by the potential:

$$
\Phi(r) = - \frac{G \, M}{\sqrt{r^2 + b^2}}
$$

### Exercise: What is the Plummer density distribution?

Derive the density profile using Poisson's equation.

```python

```

### Mass enclosed: Circular velocity in extended mass distributions?

Above, we saw how to compute the circular velocity of an orbit at radius $r$ in the point mass potential. This concept works in any spherical, central potential model, even when the mass is extended, such as in the case of the Plummer model. In more general spherical mass models, the equation for the circular velocity in terms of a derivative of the potential is still valid:

$$
v_{\rm circ}(r) = \sqrt{r \, \left|\frac{\textrm{d}\Phi}{\textrm{d}r}\right|}
$$

However, an alternate way of computing the circular velocity is to use the *mass enclosed* $M(<r)$, which specifies the mass enclosed within a radius $r$ and is defined as:

$$
M(<r) = \int^r_0 \textrm{d}\tilde{r} \, 4\pi\,\tilde{r}^2 \, \rho(\tilde{r})
$$

With the mass enclosed, the circular velocity can also be computed as:

$$
v_{\rm circ}(r) = \sqrt{\frac{G\,M(<r)}{r}}
$$

### Exercise: Implement the Plummer Potential and Compute a Circular Orbit

First, derive an expression for the mass enclosed in the Plummer model

```python

```

Next, implement a function to evaluate the circular velocity in the Plummer potential

```python

```

Now implement an acceleration function for the Plummer model:

```python
def plummer_acceleration(xyz, M, b):
    pass
```

Finally, compute a circular orbit starting at $(x, y, z) = (5, 0, 0)~\textrm{pc}$ (i.e. $r_0 = 5~\textrm{pc}$) using `leapfrog_integrate()` with the Plummer acceleration function you implemented, adopting the parameter values $M=10^7~\textrm{M}_\odot$ and $b = 8~\textrm{pc}$

```python

```

Plot the orbit in x-y and x-z projections:

```python

```

Compute a few other random bound orbits in the Plummer model (using small, random velocity offsets from the circular orbit initial conditions). Plot these orbits in x-y and x-z projections. How are these orbits different from the orbits in the point mass potential?

```python

```

## Constants and Integrals of Motion

We saw above that for generic orbits in time-independent gravitational fields, the total energy is conserved or constant at any point along the orbit. Depending on the symmetries of the mass distribution, orbits generically also have other conserved quantities, which are either *constants of motion* (functions of the phase-space coordinates and time $C_C(\boldsymbol{x}, \boldsymbol{v}, t)$ that are constant along an orbit) and/or *integrals of motion* (functions of the phase-space coordinates only $C_I(\boldsymbol{x}, \boldsymbol{v})$ that are constant along an orbit). The energy is an example of an integral of motion, but in a spherical, time-independent mass distribution, orbits also have other integrals of motion: The three components of the angular momentum vector $L_x, L_y, L_z$ (and therefore also the total angular momentum $L$).

To see this, let's compute the orbit of a random set of initial conditions in the Plummer model we used above:

```python
rng = np.random.default_rng(seed=123)

rand_x0 = rng.normal(0, 5, size=3) * u.pc
rand_v0 = rng.normal(0, 20, size=3) * u.km / u.s

rand_plummer_ts, rand_plummer_xs, rand_plummer_vs = leapfrog_integrate(
    plummer_acceleration,
    x0=rand_x0,
    v0=rand_v0,
    dt=dt,  # we defined this earlier!
    n_steps=1024,
    acc_func_args=(plummer_M, plummer_b),
)
```

```python
fig, axes = plt.subplots(
    1, 2, figsize=(10, 5), sharex=True, sharey=True, constrained_layout=True
)

# x-y
axes[0].plot(rand_plummer_xs[0], rand_plummer_xs[1])
axes[0].set_xlabel(f"$x$ [{rand_plummer_xs.unit:latex_inline}]")
axes[0].set_ylabel(f"$y$ [{rand_plummer_xs.unit:latex_inline}]")

# x-z
axes[1].plot(rand_plummer_xs[0], rand_plummer_xs[2])
axes[1].set_xlabel(f"$x$ [{rand_plummer_xs.unit:latex_inline}]")
axes[1].set_ylabel(f"$z$ [{rand_plummer_xs.unit:latex_inline}]")

for ax in axes:
    ax.set_aspect("equal")

    ax.set_xlim(-10, 10)
    ax.set_ylim(-10, 10)

fig.suptitle("Random orbit in the Plummer model", fontsize=24)
```

We can then compute the angular momentum vector at every timestep in the integration:

```python
L = np.cross(rand_plummer_xs, rand_plummer_vs, axis=0)
L = L.to(u.km / u.s * u.pc)
L.shape
```

We can then plot the value of the angular momentum in each component (and the total angular momentum) as a function of timestep:

```python
plt.figure(figsize=(6, 5))

plt.plot(plummer_ts, np.linalg.norm(L, axis=0), marker="", label="$L$ (total)")

names = ["x", "y", "z"]
for i in range(3):
    plt.plot(plummer_ts, L[i], marker="", label=f"$L_{names[i]}$")

plt.xlabel("timestep")
plt.ylabel("angular momentum")

plt.legend(loc="best", fontsize=14)

plt.tight_layout()
```

In the case of spherical, time-independent, central force fields, we saw above that energy $E$ and angular momentum $L$ (all components and the total magnitude) are integrals of motion. It turns out that the number of independent integrals of motion that an orbit has sets the degrees of freedom of an orbit in 6D phase space. That is, the reason that orbits do not fill six-dimensional volumes is because they often have many integrals of motion that reduce the dimensionality of the orbit (i.e. there are deep connections between coordinate symmetries, integrals of motion, and the geometry of orbits). However, we can create arbitrarily-many integrals of motion by considering any functions of these integrals of motion $f(E, L)$, so what is special about energy and angular momentum? Integrals of motion that reduce the number of degrees of freedom that an orbit has are called *isolating integrals of motion*: These are special integrals of motion that are "independent" (in detail, any such integrals are "in involution" such that the Poisson bracket of any pair of isolating integrals of motion must be zero; That is TMI for our purposes!). The number of isolating integrals of motion have geometric importance: Orbits exist in 6D phase space. If energy is conserved, this restricts orbits to a 5D manifold embedded in 6D phase space. In the case of spherical gravitational fields, the three components are three additional isolating integrals, which restrict the degrees of freedom to a 2D manifold: This is why generic orbits in spherical models (e.g., the Plummer model) trace out rosette patterns bound to a planar surface!

### Exercise: Why do orbits in the point mass potential not fill a 2D space, but instead trace out a 1D curve?

I.e. can you guess how many isolating integrals of motion there are for orbits in a point mass potential?

Answer here:

## A Quick Intro to Gala: Fast Orbit Integration

The point mass and Plummer potentials are just two examples of spherical, central mass models. A number of other spherical models appear often in Galactic dynamics contexts, most of which are just potential models resulting from spherical density distributions that have different power law exponents.

To demonstrate these other mass models, so that we do not have to implement all of the models here, we are going to switch from using our own home-brewed integration and potential definitions to using the [Gala](http://gala.adrian.pw/) Python package. Gala is maintained by me (Adrian) and was born out of my need for having fast tools for numerical orbit integration in Python. (For a quick intro to Gala, see the examples below, or the [Getting Started guide in the Gala documentation](http://gala.adrian.pw/en/latest/getting_started.html)) To start with, let's import the necessary Python subpackages:

```python
import gala.dynamics as gd
import gala.integrate as gi
import gala.potential as gp
from gala.units import galactic
```

As an initial demonstration, let's re-compute the same random orbit we computed in the Plummer model just above: The position and velocity initial conditions are stored in `rand_x0` and `rand_v0`, and the potential parameters are available as `plummer_M` and `plummer_b`.

The first thing we will do is define a Gala `PlummerPotential` object to represent the potential we are computing the orbit in:

```python
plummer_pot = gp.PlummerPotential(m=plummer_M, b=plummer_b, units=galactic)
plummer_pot
```

Here, we have to specify a unit system along with the potential parameters, which was also imported from Gala: In Gala, the `galactic` unit system is defined as (kpc, Myr, solar masses, radians):

```python
galactic
```

This potential object already knows how to compute, e.g., the energy, acceleration, or density for the Plummer mass model. For example, to compute the potential energy at some random position:

```python
some_position = [5.5, -0.14, 3.4] * u.pc
plummer_pot.energy(some_position)
```

or the acceleration and density:

```python
print(plummer_pot.acceleration(some_position))
print(plummer_pot.density(some_position))
```

To compute an orbit with this potential object, we now have to define a new object to represent the initial conditions. We have the values we will need saved in the `rand_x0` and `rand_v0` variables, so we just need to pass these in to the `PhaseSpacePosition` class to create a Gala object to contain the initial condition data:

```python
w0 = gd.PhaseSpacePosition(pos=rand_x0, vel=rand_v0)
```

To compute the orbit, we can use the `.integrate_orbit()` method on any Gala potential object:

```python
dt.to(u.Myr)  # defined up above!
```

```python
gala_plummer_orbit = plummer_pot.integrate_orbit(w0, dt=dt, n_steps=1_000)
```

The object returned from this orbit integration is a Gala `Orbit` object, which contains the position, velocity, and time at each timestep, but also has a number of convenience methods that enable various things, like making a quick plot of the different projections of the orbit:

```python
_ = gala_plummer_orbit.plot(units=u.pc)
```

Or computing the angular momentum at every timestep:

```python
gala_L = gala_plummer_orbit.angular_momentum()
gala_L.shape
```

Or even other quantities like the pericenter (smallest radius an orbit reaches), apocenter (largest radius an orbit reaches), or eccentricity (defined as $e = \frac{r_{\rm apo} - r_{\rm peri}}{r_{\rm apo} + r_{\rm peri}}$) of an orbit:

```python
gala_plummer_orbit.pericenter().to(u.pc)
```

```python
gala_plummer_orbit.apocenter().to(u.pc)
```

```python
gala_plummer_orbit.eccentricity()
```

The default integrator used by Gala is Leapfrog integration, but you can change the integrator to, e.g., an 8th order Runge–Kutta scheme (the Dormand–Prince 853 method) by passing this integrator to `integrate_orbit()`

```python
gala_plummer_orbit_dop853 = plummer_pot.integrate_orbit(
    w0, dt=dt, n_steps=1_000, Integrator=gi.DOPRI853Integrator
)
```

### Exercise: Compute the energy of the DOP853 orbit at all timesteps

*Hint: there is a useful method on `Orbit` objects for doing this!*

```python

```

## Common Spherical Potential Models: Isochrone Potential

The Isochrone potential is a special mass model in which the potential is given by

$$
\begin{align}
\Phi(r) &= - \frac{G \, M}{\sqrt{r^2 + b^2} + b}\\
&= - \frac{G \, M}{a + b}\\
a &\equiv \sqrt{r^2 + b^2}
\end{align}
$$

where $M$ is the total mass and $b$ is a scale radius parameter. The corresponding density distribution is

$$
\rho(r) = M \, \left( \frac{3\,(a+b)\,a^2 - r^2\,(3a + b)}{4\pi\,(a+b)^3 \, a^3} \right) \quad .
$$

This density function arises from softening a point mass: Even at $r=0$, this model has finite density. Well within the scale radius ($r \ll b$), the density is approximately constant:

$$
\rho_{r \ll b} \approx \frac{3 M}{16\pi \, b^3}
$$

Recalling Poisson's equation, this implies that the potential in this region scales with $r^2$ (note: not $1/r^2$!): This is the potential of a uniform density (i.e. "homogeneous") sphere. The potential of a uniform density sphere is in some sense an extremum for us to consider, because most relevant mass models will have decreasing density with increasing radius.

At the other limit in the Isochrone potential, i.e. at large radius, $r \gg b$ and $a \gg b$, so to first order, the potential is:

$$
\Phi_{r\gg b} \approx - \frac{G \, M}{r}\\
$$

which you may recall is the potential of a point mass!

The Isochrone model is therefore a special model because it can "interpolate" between the potential of a uniform density sphere at small radius, and the potential of a point mass at large radius. At intermediate radii, depending on the setting of $b$, the isochrone model can be used as a good *local* approximation to many other potential models. This is useful because it is often more efficient to compute an orbit in the Isochrone model over other potentials.

### Example: Compute orbits in different regimes of the Isochrone potential

As a demonstration, we will compute 3 orbits in the Isochrone potential: One at small radius ($r\ll b$), one at intermediate radius ($r \sim b$), and one at large radius ($r \gg b$). We will use Gala to to the orbit integration:

```python
iso_b = 10 * u.kpc
isochrone_pot = gp.IsochronePotential(m=1e12 * u.Msun, b=iso_b, units=galactic)
```

Here, we will define 3 sets of initial conditions for the three orbits and store these all in an array-valued `PhaseSpacePosition` object:

```python
iso_x0 = ([[0.05, 0, 0], [1, 0, 0], [1000, 0, 0]] * iso_b).T

vcirc = isochrone_pot.circular_velocity(iso_x0)
# The line below uses a Numpy paradigm called "array broadcasting"
# to fill out the 3 x 3 array of initial velocities for the 3 orbits
# https://numpy.org/doc/stable/user/basics.broadcasting.html
iso_v0 = vcirc[np.newaxis] * np.array([0.5, 1, 0])[:, np.newaxis]

iso_w0 = gd.PhaseSpacePosition(pos=iso_x0, vel=iso_v0)
iso_w0.shape
```

We could then use `Potential.integrate_orbit()` to compute the three orbits simultaneously. However, here we expect the orbits to have very different orbital periods. We can get a crude estimate of the expected orbital period assuming circular orbits:

```python
iso_periods = (2 * np.pi * iso_w0.physicsspherical.r / vcirc).decompose(galactic)
iso_periods
```

We can see that the expected orbital periods span about 2 orders of magnitude, so we would waste a lot of computational time if we used the same (small) timestep for all three orbits. Instead, we will use a fraction of the period as our timestep, and we will compute the orbits separately:

```python
iso_orbits = []
for i in range(3):
    dt = iso_periods[i] / 256  # 256 timesteps per period
    iso_orbit = isochrone_pot.integrate_orbit(iso_w0[i], dt=dt, n_steps=4096)
    iso_orbits.append(iso_orbit)
```

```python
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
for ax, orbit in zip(axes, iso_orbits):
    orbit.plot(["x", "y"], axes=[ax])
    lim = 1.2 * orbit.apocenter().value
    ax.set_xlim(-lim, lim)
    ax.set_ylim(-lim, lim)

axes[0].set_title(r"$r \ll b$")
axes[1].set_title(r"$r \sim b$")
axes[2].set_title(r"$r \gg b$")

fig.tight_layout()
```

Note that, as we expected from the limits discussed above, the three regimes of the Isochrone potential yield three very different orbital geometries. In the inner, uniform density sphere region (left panel), the orbit is an ellipse with the center of the ellipse at the origin. In the intermediate regime (middle panel), the orbit forms a rosette pattern. In the outer, point-mass region (right panel), the orbit is close to an ellipse with the origin at one focus (like a Keplerian orbit!). Why does the orbit shape change so much in different parts of the Isochrone potential?

*Aside:* The ability of the Isochrone potential to locally approximate other spherical mass models is also useful because the Isochrone potential is one of the only potential models in which one can analytically perform a coordinate transformation to a set of coordinates known as *action-angle coordinates*. We will not cover action-angle coordinates in much detail here, but they may come up in your research projects or in other contexts during the summer school. Briefly, action-angle coordinates — often represented as $(\boldsymbol{J}, \boldsymbol{\theta}$ — are a special set of canonical coordinates in which the dynamics are simple: The momentum coordinates are the actions $\boldsymbol{J}$, and these are isolating integrals of motion (i.e. they are constant!), and the position coordinates are the angles $\boldsymbol{\theta}$, which increase linearly with time.

## Orbital Frequencies

What controls whether an orbit closes on itself and fills a 1D curve instead of filling a 2D, 3D, etc. volume? One way to understand this aspect of orbital geometry is to consider its *orbital frequencies*. For a generic planar orbit, as we are considering here in spherical potentials, orbits have two fundamental orbital frequencies: One is related to the frequency with which the orbit circulates around (the azimuthal frequency $\Omega_\phi$), and the other is related to the frequency with which the orbit oscillates in and out in radius (the radial frequency $\Omega_r$). In general spherical potentials, these frequencies will typically be irrational numbers, and they are set by the local slope of the density profile near an orbit. However, in certain mass models, or certain locations in general spherical models, these two frequencies can be *resonant*: Here, two frequencies are resonant if we can find two integers $n$ and $m$ such that $n\,\Omega_\phi + m\,\Omega_r = 0$. In other words, the ratio of the orbital frequencies is a ratio of integers: $-\frac{\Omega_\phi}{\Omega_r} = \frac{m}{n}$. Orbits that satisfy a resonance relation (a particular setting of $n$ and $m$) will close on themselves: For example, if $n=2$ and $m=1$, every revolution around in $\phi$, there will be exactly 2 oscillations in $r$, bringing the orbital trajectory back to exactly the same place in phase space. Note that if the ratio was irrational, the orbit would not close because it would take an arbitrarily long time to come back to exactly the same place.

Let's estimate the orbital frequencies of the three orbits above and compute the frequency ratios to see if this gives us a better understanding of why the orbits have different morphologies. A quick and hacky way of estimating the frequencies is to find all of the local maxima in the radius $r$ and azimuth $\phi$, difference the times, and take the mean. Thankfully, `scipy` provides utilities for identifying local maxima and minima, so the procedure to estimate the orbital frequencies (periods) will be:
1. convert the orbits to spherical coordinates,
2. use the `scipy.signal.argrelmax` function to find the indices of all local maxima of $r$ and $\phi$,
3. use the returned indices to get the times corresponding to the maxima,
4. difference the subsequent times and take the mean of differences to estimate the period

```python
orbit_titles = [r"$r \ll b$", r"$r \sim b$", r"$r \gg b$"]
for orbit, name in zip(iso_orbits, orbit_titles):
    # 1. Convert to spherical coordinates:
    sph = orbit.physicsspherical

    # 2. Compute the indices of local maxima for r and phi:
    (r_idx,) = argrelmax(sph.r.value)
    (phi_idx,) = argrelmax(sph.phi.value)

    # 3. Retrieve the corresponding times:
    r_times = orbit.t[r_idx]
    phi_times = orbit.t[phi_idx]

    # 4. Difference the times and take the mean to estimate
    # the orbital periods in each coordinate:
    P_r = np.mean(np.diff(r_times))
    P_phi = np.mean(np.diff(phi_times))

    # Convert to frequencies:
    Omega_r = 2 * np.pi * u.rad / P_r
    Omega_phi = 2 * np.pi * u.rad / P_phi

    # Print the frequency ratio:
    display(Math(f"{name}  :  $\Omega_\phi / \Omega_r$ = {Omega_phi / Omega_r:.3f}"))
```

For the inner (~uniform density; left) orbit, the frequency ratio is very close to 1/2, meaning that for every revolution in $\phi$, there is very close to 2 oscillations in $r$! This explains the fact that the orbital shape looks like an ellipse, but with the center of the ellipse at the origin.

For the intermediate orbit (middle), the frequency ratio looks irrational, implying that the orbit should fill a 2D region. This explains the rosette pattern we see in the middle panel of the figure above. The fact that this frequency ratio is between 0.5 and 1 means that the radial frequency is always longer (i.e. the period is shorter), so orbits tend to complete between 1 and 2 radial oscillations for every azimuthal revolution.

For the outer orbit (~point mass; right), the frequency ratio is very close to 1, meaning that for every revolution in $\phi$, there is very close to 1 oscillation in $r$. This explains the elliptical shape and property of Keplerian orbits: The focus of the ellipse is at the origin because when the orbit reaches phase = 1/2 in azimuth, it also reaches phase = 1/2 in radius.

Orbits in most common spherical models relevant for Galactic dynamics have irrational azimuthal-to-radial frequency ratios between 0.5 and 1, meaning that generic orbits typically look like rosettes.

## Common Spherical Potential Models: Power-law Models

Many other commonly-used spherical potentials are derived from either "single power" or "two power" density distributions.

### Single-power Models

Single-power density distributions have a form like:

$$
\rho(r) = \frac{\rho_0}{(r/r_0)^\alpha}
$$

Which implies a mass enclosed profile that depends on the power-law index $\alpha$:

$$
M(<r) \propto r^{3-\alpha}
$$

If you recall the expressions for the circular velocity

$$
v_\textrm{circ}(r) = \sqrt{\frac{G\,M(<r)}{r}}
$$

this implies that power-law density profiles have circular velocity curves:

$$
v^2_{\textrm{circ}, \textrm{SPL}}(r) \propto \frac{r^{3-\alpha}}{r}\\
v_{\textrm{circ}, \textrm{SPL}}(r) \propto r^{1-\alpha/2}
$$

When $\alpha = 2$, this implies that the circular velocity is independent of radius! This case corresponds to a potential that has the form:

$$
\Phi_{\alpha=2}(r) = \frac{1}{2}v_c^2 \, \ln\left(\frac{r}{r_0}\right)
$$

This *Logarithmic potential* is sometimes used as a simple model for galaxies, because they are observed to have approximately flat rotation curves (implying a close-to-constant circular velocity).

#### Example: Comparing the circular velocity curves of several potential models

Here, we will use Gala to compute and plot the circular velocity curves of a few of the potential models we have introduced: point mass, Plummer, Isochrone, and Logarithmic

```python
potentials = {
    "Point mass": gp.KeplerPotential(m=1),
    "Plummer": gp.PlummerPotential(m=1, b=1),
    "Isochrone": gp.IsochronePotential(m=1, b=1),
    "Logarithmic": gp.LogarithmicPotential(v_c=1, r_h=1),
}

r_grid = np.linspace(0, 10, 1024)
xyz_grid = np.zeros((3, len(r_grid)))
xyz_grid[0] = r_grid

fig, ax = plt.subplots(figsize=(6, 4))
for name, pot in potentials.items():
    ax.plot(r_grid, pot.circular_velocity(xyz_grid), marker="", lw=2, label=name)

ax.set_xlim(0, 10)
ax.set_ylim(0, 2)

ax.legend(loc="upper right", fontsize=12)

ax.set_xlabel("radius")
ax.set_ylabel("circular velocity")
```

### Two-power Models

Two-power density distributions have a form like:

$$
\rho(r) = \frac{\rho_0}{(r/r_0)^\alpha \, (1 + r/r_0)^{\beta - \alpha}}
$$

Many useful potential models are found from setting particular values of $\beta$ and $\alpha$:

* NFW model (Navarro–Frenk–White): $(\alpha, \beta) = (1, 3)$
* Hernquist model: $(\alpha, \beta) = (1, 4)$
* Jaffe model: $(\alpha, \beta) = (2, 4)$
* Dehnen models: $\beta = 4$

Let's compare the potential, density, mass enclosed, and circular velocity curves of all of the extended mass models we have introduced so far:

```python
potentials = {
    "Plummer": gp.PlummerPotential(m=1, b=1),
    "Isochrone": gp.IsochronePotential(m=1, b=1),
    "Logarithmic": gp.LogarithmicPotential(v_c=1, r_h=1),
    "NFW": gp.NFWPotential(m=1, r_s=1),
    "Hernquist": gp.HernquistPotential(m=1, c=1),
    "Jaffe": gp.JaffePotential(m=1, c=1),
}

r_grid = np.linspace(0, 10, 1024)
xyz_grid = np.zeros((3, len(r_grid)))
xyz_grid[0] = r_grid

for func_name, ylim in zip(
    ["energy", "density", "mass_enclosed", "circular_velocity"],
    [(-2, 2), (1e-8, 100), (-0.5, 3), (-0.1, 1.1)],
):
    fig, ax = plt.subplots(figsize=(6, 4))
    for name, pot in potentials.items():
        ax.plot(r_grid, getattr(pot, func_name)(xyz_grid), marker="", lw=2, label=name)

    ax.set_xlim(0, 10)
    ax.set_ylim(ylim)

    if func_name == "density":
        ax.set_yscale("log")

    ax.legend(loc="best", fontsize=10)

    ax.set_xlabel("radius")
    ax.set_ylabel(func_name)
```

All of the spherical mass models above have utility in different contexts: Some have density profiles that are similar to observed light distributions in galaxies, some have circular velocity curves that match observed rotation curves in galaxies, and some are just fast to compute! There are many other possible choices for mass models that have, for example, truncated density distributions, or more complex radial dependencies, but often at the expense of simplicity: For example, some will require numerically solving for the potential or gradient, which introduces some overhead over the (mostly) analytic models discussed above. You may therefore encounter cases where people use sums or combinations of the models given above as a way of approximating more complex mass distributions, because potentials are additive.

So far, we have only discussed spherical mass distributions and the properties of orbits in spherical models. In the next tutorial notebook, we will introduce some common axisymmetric and triaxial density distributions and we will explore what happens to orbits when we consider more complex mass distributions.


```python
from IPython.display import Image
```

# 2D Turbulent Hot Free Jet

## Literature

[**"Physical and computational aspects of convective heat transfer"**](http://link.springer.com/book/10.1007%2F978-1-4612-3918-5)

T. CEBECI, P. BRADSHAW, Springer 1984

## Equations for the 2D hot free jet

Applying assumptions from [Prandtl's boundary layer](https://en.wikipedia.org/wiki/Boundary_layer) theorem, the [Navier-Stokes equations](https://en.wikipedia.org/wiki/Navier%E2%80%93Stokes_equations) can simplified for the 2D turbulent free jet. The follwoing set of equations depict the continuity, momentum and energy equations, respectively.

\begin{align}
\frac{\partial u}{\partial x} + \frac{\partial v}{\partial y} & = 0 \\
u\frac{\partial u}{\partial x} + v\frac{\partial u}{\partial y} & = \frac{1}{\rho}\frac{\partial \tau}{\partial y} \\
u\frac{\partial u}{\partial x} + v\frac{\partial u}{\partial y} & = - \frac{1}{\rho c_P}\frac{\partial \dot q}{\partial y}
\end{align}

With

\begin{equation}
\tau = \mu \frac{\partial u}{\partial y}, \qquad \tau = \mu \frac{\partial u}{\partial y} - \rho \overline{u'v'}
\end{equation}

for the laminar and turbulent case, and similarly we have

\begin{equation}
\dot q = -\lambda \frac{\partial T}{\partial y}, \qquad \dot q = -\lambda \frac{\partial T}{\partial y} + \rho c_P \overline{T'v'}
\end{equation}

for the laminar and turbulent cases of the energy equation. A proper turbulence model has to be selected for the *Reynolds' stresses* and *turbulent diffusion*. In this case the **Cebeci-Smith** model, derived from the Schlichting formula which is based on the mixing length is applied.

Analogy considerations to the molecular viscosity lead to

\begin{equation}
\mu_t = \frac{- \rho \overline{u'v'}} {\partial u/\partial y} \qquad \text{so that} \qquad \mu_{eff} = \mu + \mu_t
\end{equation}

and for the turbulent diffussivity

\begin{equation}
\kappa_t = \frac{-\rho c_P \overline{T'v'}} {\partial T/\partial y} \qquad \text{so that} \qquad \kappa_{eff} = \kappa + \kappa_t  \qquad \text{with}  \qquad \kappa = \frac{\lambda}{\rho c_P} \qquad \text{being the } \textit {thermal diffusivity}.
\end{equation}

With

\begin{equation}
\epsilon_m = \frac{\mu_t}{\rho} \quad \text{and} \quad \epsilon_h = \frac{\kappa_t}{\rho c_P},\quad \text {we define the } \textit{eddy kinematic viscosity} \text{ and } \textit{eddy diffusivity of heat} \text{, respectively}.
\end{equation}

By analogy with the molecular Prandtl number:

\begin{equation}
Pr \equiv \frac{\nu}{\lambda} \equiv \frac{\mu c_P}{\lambda}
\end{equation}

 a *turbulent Prandtl* number can be defined

\begin{equation}
Pr_t \equiv \frac{\epsilon_m}{\epsilon_h}  \equiv \frac{\epsilon_m^+}{\epsilon_h^+} \equiv \frac{\rho \overline{u'v'} / (\partial u/\partial y)}{\overline{T'v'} / (\partial T/\partial y)}
\end{equation}

with

\begin{equation}
\epsilon_m^+ \equiv \frac{\epsilon_m}{\nu} \quad \text{and} \quad \epsilon_h^+ \equiv \frac{\epsilon_h}{\nu}.
\end{equation}

 Applying the **Falkner-Skan** transformation:

\begin{equation}
\eta \equiv \frac{y}{\delta} = \sqrt{\frac{u_0}{\nu L}}\frac{y}{3\xi^{2/3}}, \qquad \xi=\frac{x}{L}, \qquad \psi = \sqrt{u_0 \nu L}\xi^{1/3}f(\xi, \eta), \quad u = \frac{\partial \psi}{\partial y} = \frac{\partial \psi}{\partial \eta}\frac{\partial \eta}{\partial y}, \quad v = -\frac{\partial \psi}{\partial x} = -\frac{\partial \psi}{\partial \eta}\frac{\partial \eta}{\partial x}
\end{equation}

\begin{equation}
u = \frac{1}{3} \frac{u_0 f'}{\xi^{1/3}} \quad \text{with} \quad f' \equiv \frac{\partial f}{\partial \eta}
\end{equation}

the following equations for the 2D laminar heated free jet can be derived:

\begin{align}
f'''+(f')^2 +ff'' &= 3 \xi \left( f' \frac{\partial f'}{\partial \xi} - f'' \frac{\partial f}{\partial \xi}\right) \\
\frac{1}{Pr}g'' + (fg)' &= 3 \xi \left( f' \frac{\partial g}{\partial \xi} - g' \frac{\partial f}{\partial \xi}\right)
\end{align}

The boundary conditions are:

\begin{align}
\eta = 0&: \quad f = f'' = 0, \quad g' = 0 \\
\eta = \eta_e&: \quad f' = 0, \quad g = 0
\end{align}

## Keller's BOX method

The BOX-method of H.B. Keller is a second order accurate discretization scheme for parabolic partial differential equations. Since the boundary layer equations are of parabolic type the scheme is well suited.

The BOX-method consists of four steps:
1. Reduction of the order of the PDE by substitution to get a first order equation system
2. Use central differences for the discretization of the first order equations
3. Linearize the resulting algebraic equations (in case they are non-linear)
4. Solve the linear system

The order recuction is done by substitution which typically leads to a set of ODEs and one PDE. The ODEs are discretized at the new, unknown position (e.g. space or time), whereas the first oder PDE is discretized halfway between known and unknown quantities.

In this case, steps 3 and 4 are treated in a different way and a solver for a system of nonlinear equations is employed instead.

1. Reduction of the order of the PDE by substitution to get a first order equation system
2. Use central differences for the discretization of the first order equations
3. Rewrite linear and nonlinear discretized equations in the form F(x) = 0 and solve the system using SCIPY.

The solver applied is the **fsolve** solver from **scipy.optimize** and is dedicated to large systems of equations forming sparse matrices.

### Step 1 - Order Reduction
Applying the order reduction of the BOX-method to the equations for the free jet and also introducing turbulence yields the following set of equations:

\begin{align}
f' &= u(\xi, \eta) \\
u' &= v(\xi, \eta) \\
g' &= p(\xi, \eta) \\
(bv)'+u^2 +fv &= 3 \xi \left( u \frac{\partial u}{\partial \xi} - v \frac{\partial f}{\partial \xi}\right) \\
(ep)' + fp &= 3 \xi \left( u \frac{\partial g}{\partial \xi} - p \frac{\partial f}{\partial \xi}\right)
\end{align}

which is a set of 3 ODEs and 2 PDEs, where

\begin{equation}
b = 1 + \epsilon_m^+  \quad \text{and} \quad e = \frac{1}{Pr} + \frac{\epsilon_m^+}{Pr_t}
\end{equation}

The boundary conditions after the substitution are:

\begin{align}
\eta = 0&: \quad f = v = 0, \quad p = 0 \\
\eta = \eta_e&: \quad u = 0, \quad g = 0
\end{align}

### Step 2 - Discretization

The figures below depict the transformed computational mesh at the nozzle exit and the nomenclature of the *Keller BOX* and associated discretization points for the ODEs and PDEs. In the lower sketch the solution propagation of the PDE is from left to right which corresponds to the flow direction and the parabolic nature of the flow type (downstream disturbances do not affect the upstream solution).

```python
import os
path = 'C:\Dropbox\JET\Documentation'
_file = 'Nozzle_Meshgrid.png'
filename = os.path.join(path, _file)
Image(filename=filename, width=600)
```

```python
_file = 'Keller_BOX.png'
filename = os.path.join(path, _file)
Image(filename=filename, width=600)
```

#### ODE Discretization

\begin{equation}
f' = \frac{df}{d\eta} = u
\end{equation}

Using central differences to dicretize this ODE at the unknown location (see *ODE discretization point* in above figure) results in:

\begin{equation}
\frac{1}{h_j}(f_{j}^{n}-f_{j-1}^{n}) = \frac{1}{2}(u_{j}^{n}+u_{j-1}^{n})
\end{equation}

and likewise for the remaining two ODEs:

\begin{equation}
\frac{1}{h_j}(u_{j}^{n}-u_{j-1}^{n}) = \frac{1}{2}(v_{j}^{n}+v_{j-1}^{n})
\end{equation}

\begin{equation}
\frac{1}{h_j}(g_{j}^{n}-g_{j-1}^{n}) = \frac{1}{2}(p_{j}^{n}+p_{j-1}^{n})
\end{equation}

####PDE Discretization

#####Discretization of the Momentum Equation

\begin{equation}
(bv)'+u^2 +fv = 3 \xi \left( u \frac{\partial u}{\partial \xi} - v \frac{\partial f}{\partial \xi}\right)
\end{equation}

The PDE is again discretized using central differences, but now at a point halfway between known and unknown quantities in the center of the "Keller BOX" (see *PDE disecretization point* in above figure).

The discretization is done one by one for each individual term. For the first term of the momentum equation it is done in a quite verbose way in order to show the process in every detail. In the last line the variables are already separated in terms of known (superscript $n-1$) and unknown (superscript $n$) positions. This is needed later for grouping the coefficients of the unknowns as well as the known terms for the matrix-vector solution procedure ($\matrix{A}\vec{x}=\vec{b}$) for the linear equation system.

\begin{align}
(bv)' = \frac{\partial (bv)}{\partial \eta} &\approx \frac{1}{h_j}\left [\left(bv\right)_{j}^{n-1/2}-\left(bv\right)_{j-1}^{n-1/2}\right ] =\\
&= \frac{1}{h_j}\left [ \frac{1}{2}\left((bv)_{j}^{n} + (bv)_{j}^{n-1}\right) - \frac{1}{2}\left((bv)_{j-1}^{n} + (bv)_{j-1}^{n-1}\right)\right] = \\
&= \frac{1}{h_j}\left [ \frac{1}{2}(b_{j}^{n}v_{j}^{n} + b_{j}^{n-1}v_{j}^{n-1}) - \frac{1}{2}(b_{j-1}^{n}v_{j-1}^{n} + b_{j-1}^{n-1}v_{j-1}^{n-1})\right] = \\
&= \frac{1}{2 h_j} \left(b_{j}^{n}v_{j}^{n} - b_{j-1}^{n}v_{j-1}^{n}\right) + \frac{1}{2 h_j}\left(b_{j}^{n-1}v_{j}^{n-1} - b_{j-1}^{n-1}v_{j-1}^{n-1}\right) \\
\quad \\
u^2 &\approx \left (u^2 \right )_{j-1/2}^{n-1/2} = \\
& = \frac{1}{2} \left[\left ( u^2 \right)_{j-1/2}^{n} + \left(u^2\right)_{j-1/2}^{n-1}\right] \\
& = \frac{1}{4}\left[\left ( u^2 \right)_{j}^{n} + \left ( u^2 \right)_{j-1}^{n}\right] + \frac{1}{4}\left[\left ( u^2 \right)_{j}^{n-1} + \left ( u^2 \right)_{j-1}^{n-1}\right]\\
\quad \\
fv &\approx \left (fv \right )_{j-1/2}^{n-1/2} = \\
& = \frac{1}{2} \left[ \left(fv\right)_{j-1/2}^{n} + \left(fv\right)_{j-1/2}^{n-1}\right] \\
& = \frac{1}{4}\left(f_{j}^{n}v_{j}^{n} + f_{j-1}^{n}v_{j-1}^{n}\right) + \frac{1}{4}\left ( f_{j}^{n-1}v_{j}^{n-1} + f_{j-1}^{n-1}v_{j-1}^{n-1}\right) \\
\quad \\
3\xi &= 3\xi^{n-1/2} = \frac{3}{2}\left(\xi^n+\xi^{n-1}\right)
\quad \\
u &\approx u_{j-1/2}^{n-1/2} = \\
& = \frac{1}{2} \left(u_{j-1/2}^{n} + u_{j-1/2}^{n-1}\right) \\
& = \frac{1}{4}\left(u_{j}^{n} + u_{j-1}^{n}\right) + \frac{1}{4}\left(u_{j}^{n-1} + u_{j-1}^{n-1}\right)\\
\quad \\
\frac{\partial u}{\partial \xi} &\approx \frac{1}{k_n}\left (u_{j-1/2}^{n}-u_{j-1/2}^{n-1}\right ) =\\
&= \frac{1}{2 k_n}\left (u_{j}^{n}+u_{j-1}^{n}\right) - \frac{1}{2 k_n}\left (u_{j}^{n-1}+u_{j-1}^{n-1}\right) = \\
\quad \\
v &\approx v_{j-1/2}^{n-1/2} = \\
& = \frac{1}{2} \left(v_{j-1/2}^{n} + v_{j-1/2}^{n-1}\right) \\
& = \frac{1}{4}\left(v_{j}^{n} + v_{j-1}^{n}\right) + \frac{1}{4}\left(v_{j}^{n-1} + v_{j-1}^{n-1}\right)\\
\quad \\
\frac{\partial f}{\partial \xi} &\approx \frac{1}{k_n}\left (f_{j-1/2}^{n}-f_{j-1/2}^{n-1}\right ) =\\
&= \frac{1}{2 k_n}\left (f_{j}^{n}+f_{j-1}^{n}\right) - \frac{1}{2 k_n}\left (f_{j}^{n-1}+f_{j-1}^{n-1}\right)
\end{align}

Keeping the $( )_{j-1/2}$ terms, the complete discretized reduced order momentum equation now is:

\begin{equation}
\frac{1}{2 h_j} \left(b_{j}^{n}v_{j}^{n} - b_{j-1}^{n}v_{j-1}^{n}\right) + \frac{1}{2 h_j}\left(b_{j}^{n-1}v_{j}^{n-1} - b_{j-1}^{n-1}v_{j-1}^{n-1}\right) + \frac{1}{2} \left[\left ( u^2 \right)_{j-1/2}^{n} + \left(u^2\right)_{j-1/2}^{n-1}\right] + \frac{1}{2} \left[\left ( fv \right)_{j-1/2}^{n} + \left(fv\right)_{j-1/2}^{n-1}\right] = \frac{3}{2 k_n}\left(\xi^n+\xi^{n-1}\right)\left \{ \frac{1}{2} \left(u_{j-1/2}^{n} + u_{j-1/2}^{n-1}\right)\left (u_{j-1/2}^{n}-u_{j-1/2}^{n-1}\right) + \frac{1}{2} \left(v_{j-1/2}^{n} + v_{j-1/2}^{n-1}\right) \left(f_{j-1/2}^{n}-f_{j-1/2}^{n-1}\right )   \right\}
\end{equation}

Multiply by 2 and regroup LHS by unknowns and knowns:

\begin{equation}
\frac{1}{h_j} \left(b_{j}^{n}v_{j}^{n} - b_{j-1}^{n}v_{j-1}^{n}\right) + \left ( u^2 \right)_{j-1/2}^{n} + \left ( fv \right)_{j-1/2}^{n} + \frac{1}{h_j}\left(b_{j}^{n-1}v_{j}^{n-1} - b_{j-1}^{n-1}v_{j-1}^{n-1}\right)  + \left(u^2\right)_{j-1/2}^{n-1} + \left(fv\right)_{j-1/2}^{n-1} = \frac{3}{2 k_n}\left(\xi^n+\xi^{n-1}\right) \left\{ \left(u_{j-1/2}^{n} + u_{j-1/2}^{n-1}\right)\left (u_{j-1/2}^{n}-u_{j-1/2}^{n-1}\right) - \left(v_{j-1/2}^{n} + v_{j-1/2}^{n-1}\right) \left(f_{j-1/2}^{n}-f_{j-1/2}^{n-1}\right )   \right\}
\end{equation}

The prefactor on the RHS is substitued by $\alpha$:

\begin{equation}
\alpha = \frac{3}{2 k_n}\left(\xi^n+\xi^{n-1}\right) = 3 \frac{\xi^{n-1/2}}{\xi^n-\xi^{n-1}}
\end{equation}

\begin{equation}
\frac{1}{h_j} \left(b_{j}^{n}v_{j}^{n} - b_{j-1}^{n}v_{j-1}^{n}\right) + \left ( u^2 \right)_{j-1/2}^{n} + \left ( fv \right)_{j-1/2}^{n} + \frac{1}{h_j}\left(b_{j}^{n-1}v_{j}^{n-1} - b_{j-1}^{n-1}v_{j-1}^{n-1}\right)  + \left(u^2\right)_{j-1/2}^{n-1} + \left(fv\right)_{j-1/2}^{n-1} = \alpha \left\{ \left(u_{j-1/2}^{n} + u_{j-1/2}^{n-1}\right)\left (u_{j-1/2}^{n}-u_{j-1/2}^{n-1}\right) - \left(v_{j-1/2}^{n} + v_{j-1/2}^{n-1}\right) \left(f_{j-1/2}^{n}-f_{j-1/2}^{n-1}\right )   \right\}
\end{equation}

Resolving the terms on the RHS and indicating knowns and unknowns at the LHS:

\begin{equation}
\underbrace{ \frac{1}{h_j} \left(b_{j}^{n}v_{j}^{n} - b_{j-1}^{n}v_{j-1}^{n}\right) + \left ( u^2 \right)_{j-1/2}^{n} + \left ( fv \right)_{j-1/2}^{n} }_\text{unknown} + \underbrace{ \frac{1}{h_j}\left(b_{j}^{n-1}v_{j}^{n-1} - b_{j-1}^{n-1}v_{j-1}^{n-1}\right)  + \left(u^2\right)_{j-1/2}^{n-1} + \left(fv\right)_{j-1/2}^{n-1} }_\text{known} = \alpha \left\{ \left(u^2\right)_{j-1/2}^{n} - u_{j-1/2}^{n}u_{j-1/2}^{n-1}+ u_{j-1/2}^{n-1}u_{j-1/2}^{n}  - \left(u^2\right)_{j-1/2}^{n-1} - \left(fv\right)_{j-1/2}^{n} + v_{j-1/2}^{n}f_{j-1/2}^{n-1} - v_{j-1/2}^{n-1}f_{j-1/2}^{n} + \left(fv\right)_{j-1/2}^{n-1}   \right\}
\end{equation}

Futher rerarrangement and putting all unknowns to the LHS (also those, where knowns and unknowns are combined in a term) and all knowns to the RHS yields the final discretization of the reduced order momentum equation:

\begin{equation}
\frac{1}{h_j} \left(b_{j}^{n}v_{j}^{n} - b_{j-1}^{n}v_{j-1}^{n}\right) + \left(1-\alpha\right) \left( u^2 \right)_{j-1/2}^{n} + \left(1+\alpha\right) \left( fv \right)_{j-1/2}^{n} + \alpha \left( v_{j-1/2}^{n-1}f_{j-1/2}^{n}  - v_{j-1/2}^{n}f_{j-1/2}^{n-1} \right) = - \frac{1}{h_j}\left(b_{j}^{n-1}v_{j}^{n-1} - b_{j-1}^{n-1}v_{j-1}^{n-1}\right) - \left(1+\alpha\right) \left( u^2 \right)_{j-1/2}^{n-1} - \left(1-\alpha\right)\left(fv\right)_{j-1/2}^{n-1}
\end{equation}

#####Discretization of the Energy Equation

\begin{equation}
(ep)' + fp = 3 \xi \left( u \frac{\partial g}{\partial \xi} - p \frac{\partial f}{\partial \xi}\right)
\end{equation}

The discretization is done again one by one for each individual term, but now without being too verbose.

\begin{align}
(ep)' = \frac{\partial (ep)}{\partial \eta} &\approx \frac{1}{h_j}\left [\left(ep\right)_{j}^{n-1/2}-\left(ep\right)_{j-1}^{n-1/2}\right ] =\\
&= \frac{1}{2 h_j} \left(e_{j}^{n}p_{j}^{n} - e_{j-1}^{n}p_{j-1}^{n}\right) + \frac{1}{2 h_j}\left(e_{j}^{n-1}p_{j}^{n-1} - e_{j-1}^{n-1}p_{j-1}^{n-1}\right) \\
\quad \\
fp &\approx \left (fp \right )_{j-1/2}^{n-1/2} = \\
& = \frac{1}{2} \left[ \left(fp\right)_{j-1/2}^{n} + \left(fp\right)_{j-1/2}^{n-1}\right] \\
\quad \\
3\xi &= 3\xi^{n-1/2} = \frac{3}{2}\left(\xi^n+\xi^{n-1}\right)
\quad \\
u &\approx u_{j-1/2}^{n-1/2} = \\
& = \frac{1}{2} \left(u_{j-1/2}^{n} + u_{j-1/2}^{n-1}\right) \\
\quad \\
\frac{\partial g}{\partial \xi} &\approx \frac{1}{k_n}\left (g_{j-1/2}^{n}-g_{j-1/2}^{n-1}\right )\\
\quad \\
p &\approx p_{j-1/2}^{n-1/2} = \\
& = \frac{1}{2} \left(p_{j-1/2}^{n} + p_{j-1/2}^{n-1}\right) \\
\quad \\
\frac{\partial f}{\partial \xi} &\approx \frac{1}{k_n}\left (f_{j-1/2}^{n}-f_{j-1/2}^{n-1}\right )
\end{align}

Discretized energy equation:

\begin{equation}
\frac{1}{2 h_j} \left(e_{j}^{n}p_{j}^{n} - e_{j-1}^{n}p_{j-1}^{n}\right) + \frac{1}{2 h_j}\left(e_{j}^{n-1}p_{j}^{n-1} - e_{j-1}^{n-1}p_{j-1}^{n-1}\right) + \frac{1}{2} \left[ \left(fp\right)_{j-1/2}^{n} + \left(fp\right)_{j-1/2}^{n-1}\right] = \frac{3}{2 k_n}\left(\xi^n+\xi^{n-1}\right) \left[\frac{1}{2} \left(u_{j-1/2}^{n} + u_{j-1/2}^{n-1}\right) \left (g_{j-1/2}^{n}-g_{j-1/2}^{n-1}\right ) - \frac{1}{2} \left(p_{j-1/2}^{n} + p_{j-1/2}^{n-1}\right)\left (f_{j-1/2}^{n}-f_{j-1/2}^{n-1}\right ) \right]
\end{equation}

Final discretization of the reduced order energy equation with separated knowns and unknowns.

\begin{equation}
\frac{1}{h_j} \left(e_{j}^{n}p_{j}^{n} - e_{j-1}^{n}p_{j-1}^{n}\right) + \left(1+\alpha\right) \left( fp \right)_{j-1/2}^{n} - \alpha \left[\left( ug \right)_{j-1/2}^{n} + u_{j-1/2}^{n-1}g_{j-1/2}^{n} - u_{j-1/2}^{n}g_{j-1/2}^{n-1} + p_{j-1/2}^{n}f_{j-1/2}^{n-1} - p_{j-1/2}^{n-1}f_{j-1/2}^{n}  \right] = - \frac{1}{h_j}\left(e_{j}^{n-1}p_{j}^{n-1} - e_{j-1}^{n-1}p_{j-1}^{n-1}\right) - \left(1-\alpha\right)\left(fp\right)_{j-1/2}^{n-1} - \alpha \left( ug \right)_{j-1/2}^{n-1}
\end{equation}

**All discretized ODEs and PDEs together:**

\begin{equation}
\frac{1}{h_j}(f_{j}^{n}-f_{j-1}^{n}) = \frac{1}{2}(u_{j}^{n}+u_{j-1}^{n}) \\
\quad \\
\frac{1}{h_j}(u_{j}^{n}-u_{j-1}^{n}) = \frac{1}{2}(v_{j}^{n}+v_{j-1}^{n}) \\
\quad \\
\frac{1}{h_j}(g_{j}^{n}-g_{j-1}^{n}) = \frac{1}{2}(p_{j}^{n}+p_{j-1}^{n}) \\
\quad \\
\frac{1}{h_j} \left(b_{j}^{n}v_{j}^{n} - b_{j-1}^{n}v_{j-1}^{n}\right) + \left(1-\alpha\right) \left( u^2 \right)_{j-1/2}^{n} + \left(1+\alpha\right) \left( fv \right)_{j-1/2}^{n} + \alpha \left( v_{j-1/2}^{n-1}f_{j-1/2}^{n}  - v_{j-1/2}^{n}f_{j-1/2}^{n-1} \right) = - \frac{1}{h_j}\left(b_{j}^{n-1}v_{j}^{n-1} - b_{j-1}^{n-1}v_{j-1}^{n-1}\right) - \left(1+\alpha\right) \left( u^2 \right)_{j-1/2}^{n-1} - \left(1-\alpha\right)\left(fv\right)_{j-1/2}^{n-1}  \\
\quad \\
\frac{1}{h_j} \left(e_{j}^{n}p_{j}^{n} - e_{j-1}^{n}p_{j-1}^{n}\right) + \left(1+\alpha\right) \left( fp \right)_{j-1/2}^{n} - \alpha \left[\left( ug \right)_{j-1/2}^{n} + u_{j-1/2}^{n-1}g_{j-1/2}^{n} - u_{j-1/2}^{n}g_{j-1/2}^{n-1} + p_{j-1/2}^{n}f_{j-1/2}^{n-1} - p_{j-1/2}^{n-1}f_{j-1/2}^{n}  \right] = - \frac{1}{h_j}\left(e_{j}^{n-1}p_{j}^{n-1} - e_{j-1}^{n-1}p_{j-1}^{n-1}\right) - \left(1-\alpha\right)\left(fp\right)_{j-1/2}^{n-1} - \alpha \left( ug \right)_{j-1/2}^{n-1}
\end{equation}

And the boundary conditions become:

\begin{equation}
f_0 = 0, \quad v_0 = 0, \quad p_0 = 0 \\
u_J = 0, \quad g_J = 0
\end{equation}

### Step 3 - Rewrite Equations for SciPy Solver

If we simply rearrange the equations in the form $f(x)=0$, where $f$ is a vector valued function consisting of:

\begin{equation}
f_1(x_1, x_2, x_3, . . . , x_n) = 0 \\
f_2(x_1, x_2, x_3, . . . , x_n) = 0 \\
. \\
. \\
f_n(x_1, x_2, x_3, . . . , x_n) = 0
\end{equation}

and $(x_1, x_2, x_3, . . . , x_n)$ being the solution vector $x$ we can apply the solver [*scipy.optimize.fsolve*](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.fsolve.html#scipy.optimize.fsolve) from the [scipy](http://www.scipy.org/) package.

\begin{equation}
\frac{1}{h_j}(f_{j}^{n}-f_{j-1}^{n}) - \frac{1}{2}(u_{j}^{n}+u_{j-1}^{n}) = 0\\
\quad \\
\frac{1}{h_j}(u_{j}^{n}-u_{j-1}^{n}) - \frac{1}{2}(v_{j}^{n}+v_{j-1}^{n}) = 0\\
\quad \\
\frac{1}{h_j}(g_{j}^{n}-g_{j-1}^{n}) - \frac{1}{2}(p_{j}^{n}+p_{j-1}^{n}) = 0\\
\quad \\
\frac{1}{h_j} \left(b_{j}^{n}v_{j}^{n} - b_{j-1}^{n}v_{j-1}^{n}\right) + \left(1-\alpha\right) \left( u^2 \right)_{j-1/2}^{n} + \left(1+\alpha\right) \left( fv \right)_{j-1/2}^{n} + \alpha \left( v_{j-1/2}^{n-1}f_{j-1/2}^{n}  - v_{j-1/2}^{n}f_{j-1/2}^{n-1} \right) + \frac{1}{h_j}\left(b_{j}^{n-1}v_{j}^{n-1} - b_{j-1}^{n-1}v_{j-1}^{n-1}\right) + \left(1+\alpha\right) \left( u^2 \right)_{j-1/2}^{n-1} + \left(1-\alpha\right)\left(fv\right)_{j-1/2}^{n-1}  = 0 \\
\quad \\
\frac{1}{h_j} \left(e_{j}^{n}p_{j}^{n} - e_{j-1}^{n}p_{j-1}^{n}\right) + \left(1+\alpha\right) \left( fp \right)_{j-1/2}^{n} - \alpha \left[\left( ug \right)_{j-1/2}^{n} + u_{j-1/2}^{n-1}g_{j-1/2}^{n} - u_{j-1/2}^{n}g_{j-1/2}^{n-1} + p_{j-1/2}^{n}f_{j-1/2}^{n-1} - p_{j-1/2}^{n-1}f_{j-1/2}^{n}  \right] + \frac{1}{h_j}\left(e_{j}^{n-1}p_{j}^{n-1} - e_{j-1}^{n-1}p_{j-1}^{n-1}\right) + \left(1-\alpha\right)\left(fp\right)_{j-1/2}^{n-1} + \alpha \left( ug \right)_{j-1/2}^{n-1} = 0
\end{equation}

```python

```

```python

```

```python

```


```python
import numpy as np
import pandas as pd
import linearsolve as ls
import matplotlib.pyplot as plt
plt.style.use('classic')
%matplotlib inline
```

# Class 14: Prescott's Real Business Cycle Model I

In this notebook, we'll consider a centralized version of the model from pages 11-17 in Edward Prescott's article "Theory Ahead of Business Cycle Measurement in the Fall 1986 of the Federal Reserve Bank of Minneapolis' *Quarterly Review* (link to article: https://www.minneapolisfed.org/research/qr/qr1042.pdf). The model is just like the RBC model that we studying in the previous lecture, except that now we include an endogenous labor supply.

## Prescott's RBC Model with Labor

The equilibrium conditions for Prescott's RBC model with labor are:

\begin{align}
\frac{1}{C_t} & = \beta E_t \left[\frac{\alpha A_{t+1}K_{t+1}^{\alpha-1}L_{t+1}^{1-\alpha} +1-\delta }{C_{t+1}}\right]\\
\frac{\varphi}{1-L_t} & = \frac{(1-\alpha)A_tK_t^{\alpha}L_t^{-\alpha}}{C_t} \\
Y_t & = A_t K_t^{\alpha}L_t^{1-\alpha}\\
K_{t+1} & = I_t + (1-\delta) K_t\\
Y_t & = C_t + I_t\\
\log A_{t+1} & = \rho \log A_t + \epsilon_{t+1}
\end{align}

where $\epsilon_{t+1} \sim \mathcal{N}(0,\sigma^2)$.

The objective is use `linearsolve` to simulate impulse responses to a TFP shock using the following parameter values for the simulation:

| $\rho$ | $\sigma$ | $\beta$ | $\varphi$ | $\alpha$ | $\delta $ |
|--------|----------|---------|-----------|----------|-----------|
| 0.75   | 0.006    | 0.99    | 1.7317    | 0.35     |  0.025    |

The value for $\beta$ implies a steady state (annualized) real interest rate of about 4 percent:

\begin{align}
4 \cdot \left(\beta^{-1} - 1\right)  & \approx 0.04040
\end{align}

$\rho = 0.75$ and $\sigma = 0.006$ are consistent with the statistical properties of the cyclical component of TFP in the US. $\alpha$ is set so that, consistent with the long-run average of the US, the labor share of income is about 65 percent of GDP. The deprecation rate of capital is calibrated to be about 10 percent annually. Finally, $\varphi$ was chosen last to ensure that in the steady state households allocate about 33 percent of their available time to labor.

## Model Preparation

Before proceding, let's recast the model in the form required for `linearsolve`. Write the model with all variables moved to the left-hand side of the equations and dropping the expecations operator $E_t$ and the exogenous shock $\epsilon_{t+1}$:

\begin{align}
0 & = \beta\left[\frac{\alpha A_{t+1}K_{t+1}^{\alpha-1}L_{t+1}^{1-\alpha} +1-\delta }{C_{t+1}}\right] - \frac{1}{C_t}\\
0 & = \frac{(1-\alpha)A_tK_t^{\alpha}L_t^{-\alpha}}{C_t} - \frac{\varphi}{1-L_t}\\
0 & = A_t K_t^{\alpha}L_t^{1-\alpha} - Y_t\\
0 & = I_t + (1-\delta) K_t - K_{t+1}\\
0 & = C_t + I_t - Y_t\\
0 & = \rho \log A_t - \log A_{t+1}
\end{align}

Remember, capital and TFP are called *state variables* because they're $t+1$ values are predetermined. Output, consumption, and investment are called a *costate* or *control* variables. Note that the model as 5 equations in 5 endogenous variables.

## Initialization, Approximation, and Solution

The next several cells initialize the model in `linearsolve` and then approximate and solve it.

```python
# Create a variable called 'parameters' that stores the model parameter values in a Pandas Series
parameters = pd.Series()
parameters['rho'] = .75
parameters['sigma_squared'] = 0.006**2
parameters['beta'] = 0.99
parameters['phi'] = 1.7317
parameters['alpha'] = 0.35
parameters['delta'] = 0.025

# Print the model's parameters
print(parameters)
```

    rho              0.750000
    sigma_squared    0.000036
    beta             0.990000
    phi              1.731700
    alpha            0.350000
    delta            0.025000
    dtype: float64

```python
# Create variable called 'varNames' that stores the variable names in a list with state variables ordered first
varNames = ['a','k','y','c','i','l']

# Create variable called 'shockNames' that stores an exogenous shock name for each state variable.
shockNames = ['e_a','e_k']
```

```python
# Define a function that evaluates the equilibrium conditions of the model solved for zero. PROVIDED
def equilibrium_equations(variables_forward,variables_current,parameters):

    # Parameters. PROVIDED
    p = parameters

    # Current variables. PROVIDED
    cur = variables_current

    # Forward variables. PROVIDED
    fwd = variables_forward

    # Euler equation
    euler_equation = p.beta*(p.alpha*fwd.a*fwd.k**(p.alpha-1)*fwd.l**(1-p.alpha)+1-p.delta)/fwd.c - 1/cur.c

    # Labor-labor choice
    labor_leisure = (1-p.alpha)*cur.a*cur.k**p.alpha*cur.l**(-p.alpha)/cur.c - p.phi/(1-cur.l)

    # Production function
    production_function = cur.a*cur.k**p.alpha*cur.l**(1-p.alpha) - cur.y

    # Capital evolution
    capital_evolution = cur.i + (1 - p.delta)*cur.k - fwd.k

    # Market clearing
    market_clearing = cur.c+cur.i - cur.y

    # Exogenous tfp
    tfp_process = p.rho*np.log(cur.a) - np.log(fwd.a)

    # Stack equilibrium conditions into a numpy array
    return np.array([
        euler_equation,
        labor_leisure,
        production_function,
        capital_evolution,
        market_clearing,
        tfp_process
        ])
```

Next, initialize the model using `ls.model` which takes the following required arguments:

* `equations`
* `nstates`
* `varNames`
* `shockNames`
* `parameters`

```python
# Initialize the model into a variable named 'rbc_model'
rbc_model = ls.model(equations = equilibrium_equations,
                 nstates=2,
                 varNames=varNames,
                 shockNames=shockNames,
                 parameters=parameters)
```

```python
# Compute the steady state numerically using .compute_ss() method of rbc_model
guess = [1,4,1,1,1,0.5]
rbc_model.compute_ss(guess)

# Print the computed steady state
print(rbc_model.ss)
```

    a     1.000000
    k    11.465953
    y     1.149904
    c     0.863256
    i     0.286649
    l     0.333330
    dtype: float64

```python
# Find the log-linear approximation around the non-stochastic steady state and solve using .approximate_and_solve() method of rbc_model
rbc_model.approximate_and_solve()
```

## Impulse Responses

Compute a 26 period impulse responses of the model's variables to a 0.01 unit shock to TFP in period 5.

```python
# Compute impulse responses
rbc_model.impulse(T=26,t0=5,shocks=[0.01,0])

# Print the first 10 rows of the computed impulse responses to the TFP shock
print(rbc_model.irs['e_a'].head(10))
```

        e_a         a         k         y         c         i         l
    0  0.00  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000
    1  0.00  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000
    2  0.00  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000
    3  0.00  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000
    4  0.00  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000
    5  0.01  0.010000  0.000000  0.016262  0.001812  0.059778  0.009633
    6  0.00  0.007500  0.001494  0.012481  0.002194  0.043462  0.006858
    7  0.00  0.005625  0.002544  0.009632  0.002440  0.031290  0.004795
    8  0.00  0.004219  0.003262  0.007482  0.002587  0.022224  0.003264
    9  0.00  0.003164  0.003736  0.005857  0.002660  0.015485  0.002131

Construct a $2\times3$ grid of plots of simulated TFP, output, labor, consumption, investment, and capital. Be sure to multiply simulated values by 100 so that vertical axis units are in "percent deviation from steady state."

```python
# Create figure. PROVIDED
fig = plt.figure(figsize=(18,8))

# Create upper-left axis. PROVIDED
ax = fig.add_subplot(2,3,1)
ax.plot(rbc_model.irs['e_a']['a']*100,'b',lw=5,alpha=0.75)
ax.set_title('TFP')
ax.set_ylabel('% dev from steady state')
ax.set_ylim([-0.5,2])
ax.grid()

# Create upper-center axis. PROVIDED
ax = fig.add_subplot(2,3,2)
ax.plot(rbc_model.irs['e_a']['y']*100,'b',lw=5,alpha=0.75)
ax.set_title('Output')
ax.set_ylabel('% dev from steady state')
ax.set_ylim([-0.5,2])
ax.grid()

# Create upper-right axis. PROVIDED
ax = fig.add_subplot(2,3,3)
ax.plot(rbc_model.irs['e_a']['l']*100,'b',lw=5,alpha=0.75)
ax.set_title('Labor')
ax.set_ylabel('% dev from steady state')
ax.set_ylim([-0.5,2])
ax.grid()

# Create lower-left axis. PROVIDED
ax = fig.add_subplot(2,3,4)
ax.plot(rbc_model.irs['e_a']['c']*100,'b',lw=5,alpha=0.75)
ax.set_title('Consumption')
ax.set_ylabel('% dev from steady state')
ax.set_ylim([-0.1,0.4])
ax.grid()

# Create lower-center axis. PROVIDED
ax = fig.add_subplot(2,3,5)
ax.plot(rbc_model.irs['e_a']['i']*100,'b',lw=5,alpha=0.75)
ax.set_title('Investment')
ax.set_ylabel('% dev from steady state')
ax.set_ylim([-2,8])
ax.grid()

# Create lower-right axis. PROVIDED
ax = fig.add_subplot(2,3,6)
ax.plot(rbc_model.irs['e_a']['k']*100,'b',lw=5,alpha=0.75)
ax.set_title('Capital')
ax.set_ylabel('% dev from steady state')
ax.set_ylim([-0.2,0.8])
ax.grid()

fig.tight_layout()
```


# 3. Gyakorlat - 1 DoF lengő kar
2021.02.22

```python
from IPython.display import Image
Image(filename='1.png',width=500)
```

## Feladat:

A mellékelt ábrán egy lengőkar látható, ami két elhanyagolható tömegű rúdból és a hozzá csatlakozó három tömegpontból ($m_1, m_2, m_3$) áll. A lengőkar csak az $A$ pont körül tud elfordulni. A mozgást leíró általános koordináta a $\varphi$ szög, melyet a vízszintes egyenestől mérünk. A lengőkar a Föld gravitációs terében van.
A $\varphi=0$ helyzetben az egyensúlyt a $k$ merevségű rugó biztosítja.

### Adatok:
|||
|------------------|------------------|
| $m_1$ = 2 kg     | $l$ = 1 m        |
| $m_2$ = 4 kg     | $h$ = 0,5 m      |
| $m_3$ = 3 kg     | $a$ = 0,6 m      |
| $k$ = 10$^4$ N/m   | $g$ = 9,81 m/$s^2$ |

### Részfeladatok:
1. Írja fel a mozgásegyenletet az egyensúlyi pont körül végzett kis amplitúdójú rezgést feltételezve.
2. Számítsa ki a sajátkörfrekvenciát, a sajátfrekvenciát és a periódusidőt!
3. Határozza meg a mozgástörvényt, ha a kezdeti feltételek a függőleges pozícióra $y_{\text{B}}(t=0)=-0,01$ m, és a sebességre $v_{\text{B},y}(t=0)=-1$ m/s a $B$ pontban!
4. Számítsa ki a rugóerő maximális értékét a megadott kezdeti feltételek esetén!

## Megoldás
### Szabadtest ábra

```python
Image(filename='2.png',width=500)
```

### 1. Feladat
A szabadtest ábra alapján a mozgásegyenlet könnyen meghatározható a dinamika alaptörvényének segítségével

\begin{equation}
    \dot{\Pi}_{\text{Az}}=M_{\text{Az}}.
\end{equation}
Behelyettesítés után kapjuk:

\begin{equation}
    -\Theta_{\text{A}}\ddot{\varphi}=-m_1gl\cos\varphi-m_2gh\sin\varphi-m_3gh\sin\varphi+M_{\text{r}}(\varphi),
\end{equation}

ahol $M_{\text{r}}(\varphi)$ a rugerőhöz tartozó nyomatékot jelöli. A kis amplitúdók miatt a fenti egyenlet linearizálható az egyensúlyi pont körül $\cos\varphi\approx 1$ és a $\sin\varphi\approx \varphi$ közelítéseket alkalmazva. FONTOS: A nyomatékot leíró kifejezés nem lineáris, viszont kis rezgésamplitúdók esetén az alábbi kifejezéssel jól közelíthető

\begin{equation}
    F_{\text{r}}(\varphi)\approx F_{\text{r,st}}+ka\varphi,
\end{equation}

ahol $F_{\text{r,st}}$ az egynsúlyi helyzetben fellépő statikus rugóerő, a kifejezésben szereplő második tag pedig a rezgések által keltett úgynevezett dinamikus erő $F_{\text{r,dyn}}\approx ka\varphi$. A rugóerőhoz tartozó nyomaték linearizált alakja tehát a következő alakban adható meg:

\begin{equation}
    M_{\text{r}}(\varphi)\approx M_{\text{r,st}}+ka^2\varphi,
\end{equation}

ahol $M_{\text{r,st}}=F_{\text{r,st}}a$ a statikus rugóerő nyomatéka. Könnyen beltható, hogy a statikus rugóerőből származó nyomaték és a gravitációs erő nyomatéka pont kiejtik egymást, így ez a két tag elhagyható a mozgásegyenletből:
\begin{equation}
    M_{\text{r,st}}=m_1gl.
\end{equation}

```python
import sympy as sp
sp.init_printing()
from IPython.display import display
```

```python
## A mozgásegyenlet linearizált alakja:

# szimbólumok definiálása
g,l,h, m1, m2, m3, k, a, t = sp.symbols("g,l,h, m1, m2, m3, k, a, t", real=True) # feltételezzük, hogy valósak

# Célszerű először a tehetetlenségi nyomatékot is kifejezni megadott adatok segítségével:

Θ_Az = m2*h**2+m3*h**2+m1*l**2

# Az általános koordináta definiálása az idő függvényeként
t = sp.symbols("t",real=True, positive=True)
φ_t = sp.Function('φ')(t)

# A z tengelyre számított perdület derivált az A pontban
dΠ_Az = -Θ_Az*φ_t.diff(t,2)

# A z tengelyre számított nyomaték az A pontban
M_Az = -m2*g*h*φ_t+m3*g*h*φ_t+k*a**2*φ_t

# A dinamika alapegyenlete
# (nullára rendezve)

din = dΠ_Az-M_Az
din
```

```python
# Ahhoz, hogy a mozgásegyenletet a megszokott alakra
# rendezzük, ki kell fejezni a φ(t) együtthatóit.
# Az együtthatókat a coeff() függvény segítségével
# tudjuk kinyerni az adott kifejezésből.

co1=din.coeff(φ_t)
co2=din.coeff(φ_t.diff(t,2))
display(co1,co2)
```

```python
# Leosztva az egyenletet a második derivált együtthatójával,
# a megszokott alakora jutunk.

# Itt a sympy apart() függvényét használjuk,
# hogy a kifejezésben különszedjük a φ(t) és annak deriváltjait tartalmazó tagokat.
# (Ez nem kötelező lépés, csak áttekinthetőbb így.)

mozg_egy=sp.apart(din/co2,φ_t) # sp.apart() -> melyik kifejezésből, melyik változót
mozg_egy
```

### 2. Feladat

A mozgásegyenletben $\varphi (t)$ együtthatója egyenlő a sajátkörfrekvencia négyzetével ($\omega_n^2$).

```python
# A sajátkörfrekvencia négyzetét érdemes kimenteni egy változóba.
om2=mozg_egy.coeff(φ_t)
om2
```

```python
# gyökvonással adódik a sajátfrekvencia
ω_n=sp.sqrt(om2)
ω_n
```

```python
# Az adatoknak érdemes létrehozni egy behelyettesítési listát,
# amit a numerikus értékek meghatározásánál használhatunk.
adatok=[(m1,2),(m2,4),(m3,3),(k,10**4),(l,1),(h,0.5),(a,0.6),(g,9.81)]

# sajátkörfrekvencia numerikusan [rad/s]
ω_n_num=ω_n.subs(adatok)
ω_n_num.evalf(5)
```

```python
# sajátfrekvencia numerikusan [Hz]
f_n_num=ω_n_num/(2*sp.pi).evalf()
f_n_num.evalf(5)
```

```python
# lengésidő numerikusan [s]
T_n_num=((2*sp.pi)/ω_n_num).evalf()
T_n_num.evalf(4)
```

### 3. Feladat

```python
# Kezdeti értékek definiálása.

yB_0=-0.01
vBy_0=-1

# A kezdeti értékeket az általános koordinátával,
# illetve annak deriváltjával kell megadni a differenciálegyenlet megoldásáshoz.

kezdeti_ert = {φ_t.subs(t,0): -yB_0/l.subs(adatok), φ_t.diff(t).subs(t,0): -vBy_0/l.subs(adatok)}
kezdeti_ert
```

$\displaystyle \left\{ φ{\left(0 \right)} : 0.01, \  \left. \frac{d}{d t} φ{\left(t \right)} \right|_{\substack{ t=0 }} : 1\right\}$

```python
# Sympy-ban a kezdetiértékkel megadott differenciálegyenlet
# könnyen megoldható a dsolve() függvény segítségével.
φ_t_megold = sp.dsolve(mozg_egy.expand().subs(adatok),φ_t,ics=kezdeti_ert)
φ_t_megold
```

### 4: Feladat
A rugóerő a statikus és dinamikus rugóerő komponenesek segítségével felírható:

\begin{equation}
F_r(t)=F_{r,st}+ka\varphi(t).
\end{equation}

```python
# A statikus rugóerő a gravitációs erővel egyensúlyt tartó erő:
Fr_st=m1*g*l/a

# stat+din rugóerő
Fr_t=Fr_st+k*a*φ_t

# A maximális rugóerő ott lesz ahol a φ(t) függvénynek is maximam van
φ_max=sp.symbols("φ_max")
Fr_max=Fr_st+k*a*φ_max
Fr_max
```

A $\varphi(t)$ függvény maximuma könnyen meghatározható, ha a mozgásegyenlet megoldásást a következő alakra hozzuk:

\begin{equation}
    \varphi(t)=\Phi\cos(\omega_nt+\delta).
\end{equation}

Ekkor a $\Phi$ adja meg a rezgés amplitudóját, ami a maximális $\varphi$ értékkel egyenlő. Kihasználva a trigonometrikus azonosságot:

\begin{equation}
\Phi\cos(\omega_nt+\delta)=\Phi\cos(\omega_nt)\cos(\delta)-\Phi\sin(\omega_nt)\sin(\delta),
\end{equation}

majd összehasonlítva az együtthatókat a 3. feladatban a megoldásban szereplő együtthatókkal adódik

\begin{equation}
    \cos(\omega_nt):\hspace{50pt}C_1=\Phi\cos\delta,
\end{equation}
\begin{equation}
    \sin(\omega_nt):\hspace{50pt}C_2=\Phi\sin\delta.
\end{equation}

Innét az amplitúdó, illetve a $\delta$:

\begin{equation}
    \Phi=\sqrt{C_1^2+C_2^2},
\end{equation}

\begin{equation}
    \delta=-\arctan\left(\frac{C_2}{C_1}\right).
\end{equation}

```python
# A fentiek kiszámításához vesszük a mozgásegyenlet megoldását:
φ_t_megold
```

```python
# majd a jobb oldalára alkalmazzuk a `.as_coefficients_dict()` metódust.
dict_coeff = φ_t_megold.rhs.as_coefficients_dict()
dict_coeff

# Látható, hogy így az együtthatók és a kifejezések egy `dictionary`-ben
# kerültek tárolásra.
```

```python
# Szedjük ki a sin() és cos() együtthatóit ebből a `dictionary`-ből. Pontosabban a
# `dictionary` értkékeket tartalmazó listájából.
display(dict_coeff.values())

C1=list(dict_coeff.values())[0]
C2=list(dict_coeff.values())[1]
```

    dict_values([0.0322968709613617, 0.0100000000000000])

```python
Φ=sp.sqrt(C1**2+C2**2)
δ=sp.atan(C1/C2)
```

```python
# A fáziseltolás: [rad]
δ.evalf(5)
```

```python
# Az amplitúdó: [m]
Φ.evalf(5)
```

```python
# Ebben az esetben a maximális rugóerő: [N]
Fr_max.subs(adatok).subs(φ_max,Φ).evalf(5)
```

Készítette:

            Juhos-Kiss Álmos (Alkalmazott Mechanika Szakosztály)
            Takács Dénes (BME MM) kidolgozása és ábrái alapján.

            Hibák, javaslatok:
            amsz.bme@gmail.com
            csuzdi02@gmail.com
            almosjuhoskiss@gmail.com

            2021.02.22