Source code for dynasor.tools.damped_harmonic_oscillator

from typing import Optional, Union
import numpy as np
from numpy.typing import NDArray




def acf_position_dho(
    t: Union[float, NDArray[float]],
    w0: float,
    gamma: float,
    A: Optional[float] = 1.0,
) -> Union[float, NDArray[float]]:
    r"""
    Calculate the damped damped harmonic oscillator (DHO)
    autocorrelation function for the position.
    The definition of this function can be found in the `dynasor documentation
    <https://dynasor.materialsmodeling.org/get_started/theory.html#damped-harmonic-oscillator-model>`_.  # noqa

    Parameters
    ----------
    t
        Time, usually as an array.
    w0
        Natural angular frequency of the DHO.
    gamma
        Damping of DHO.
    A
        Amplitude of the DHO.
    """

    t = np.abs(t)

    if 2 * w0 > gamma:  # underdamped
        we = np.sqrt(w0**2 - gamma**2 / 4.0)
        return A * np.exp(-gamma * t / 2.0) * (
            np.cos(we * t) + 0.5 * gamma / we * np.sin(we * t))
    elif 2 * w0 < gamma:  # overdamped
        tau = 2 / gamma
        tau_S = tau / (1 + np.sqrt(1 - (w0 * tau)**2))
        tau_L = tau / (1 - np.sqrt(1 - (w0 * tau)**2))
        return A / (tau_L - tau_S) * (tau_L * np.exp(-t/tau_L) - tau_S * np.exp(-t/tau_S))
    else:
        tau = 2 / gamma
        return A * np.exp(-t/tau) * (1 + t / tau)





def acf_velocity_dho(
    t: Union[float, NDArray[float]],
    w0: float,
    gamma: float,
    A: Optional[float] = 1.0,
) -> Union[float, NDArray[float]]:
    r"""
    Calculate the damped damped harmonic oscillator (DHO)
    autocorrelation function for the velocity.
    The definition of this function can be found in the `dynasor documentation
    <https://dynasor.materialsmodeling.org/get_started/theory.html#damped-harmonic-oscillator-model>`_.  # noqa

    Parameters
    ----------
    t
        Time, usually as an array.
    w0
        Natural angular frequency of the DHO.
    gamma
        Damping of DHO.
    A
        Amplitude of the DHO.
    """

    t = np.abs(t)

    if 2 * w0 > gamma:  # underdamped
        we = np.sqrt(w0**2 - gamma**2 / 4.0)
        return A * w0**2 * np.exp(-gamma * t / 2.0) * (
            np.cos(we * t) - 0.5 * gamma / we * np.sin(we * t))
    elif 2 * w0 < gamma:  # overdamped
        tau = 2 / gamma
        tau_S = tau / (1 + np.sqrt(1 - (w0 * tau)**2))
        tau_L = tau / (1 - np.sqrt(1 - (w0 * tau)**2))
        return A / (tau_L - tau_S) * (np.exp(-t/tau_S)/tau_S - np.exp(-t/tau_L)/tau_L)
    else:
        tau = 2 / gamma
        return A * np.exp(-t/tau) * (1 - t / tau)





def psd_position_dho(
    w: Union[float, NDArray[float]],
    w0: float,
    gamma: float,
    A: Optional[float] = 1.0,
) -> Union[float, NDArray[float]]:
    r"""
    Calculate the power spectral density (PSD) function for the damped harmonic oscillator
    (DHO), i.e., the Fourier transform of the autocorrelation function) for the position.

    The definition of this function can be found in the `dynasor documentation
    <https://dynasor.materialsmodeling.org/get_started/theory.html#damped-harmonic-oscillator-model>`_.  # noqa

    Parameters
    ----------
    w
        Angular frequency, usually as an array.
    w0
        Natural angular frequency of the DHO.
    gamma
        Damping of DHO.
    A
        Amplitude of the DHO.
    """
    return 2 * w0**2 * A * gamma / ((w**2 - w0**2)**2 + (w * gamma)**2)





def psd_velocity_dho(
    w: Union[float, NDArray[float]],
    w0: float,
    gamma: float,
    A: Optional[float] = 1.0,
) -> Union[float, NDArray[float]]:
    r"""
    Calculate the power spectral density (PSD) function for the damped harmonic oscillator
    (DHO), i.e., the Fourier transform of the autocorrelation function) for the position.

    The definition of this function can be found in the `dynasor documentation
    <https://dynasor.materialsmodeling.org/get_started/theory.html#damped-harmonic-oscillator-model>`_.  # noqa

    Parameters
    ----------
    w
        Angular frequency, usually as an array.
    w0
        Natural angular frequency of the DHO.
    gamma
        Damping of DHO.
    A
        Amplitude of the DHO.
    """
    return w**2 * psd_position_dho(w, w0, gamma, A)