"""Generators to simulate wave system dynamics.  The viewer classes can be found in QtWaves."""

################################################################
# Written in 2003-2019 by Garth Zeglin <garthz@cmu.edu>

# To the extent possible under law, the author has dedicated all copyright
# and related and neighboring rights to this software to the public domain
# worldwide. This software is distributed without any warranty.

# You should have received a copy of the CC0 Public Domain Dedication along with this software.
# If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.

################################################################
import math, logging
import numpy as np

# set up logger for module
log = logging.getLogger('waves')

# filter out most logging; the default is NOTSET which passes along everything
# log.setLevel(logging.WARNING)

################################################################
[docs]class Waves1D(object):
    """Simulation of a set of one-dimensional wave systems.

    :param channels: number of wave systems to simulate
    :param samples: number of sample points in each system
    """

    def __init__(self, channels=5, samples=20):

        # scalar constants
        self.W = channels                       # Number of waves.
        self.N = samples                        # Number of sample points for each wave.

        self.k = 200.0*np.ones((self.W))         # Spring rate for node coupling for each wave.
        self.b =   0.5*np.ones((self.W))          # Node damping factor for each wave.
        self.k = self.k / (1 + np.arange(self.W))

        self.period    =  8.0*np.ones((self.W))              # Period of the default excitation for each wave.
        self.duty      = 0.25*np.ones((self.W))              # Duty cycle of the default excitation for each wave.
        self.magnitude =  2.0*np.ones((self.W))              # Magnitude of the default excitation for each wave.
        self.qE        =     np.zeros((self.W))              # Position of the driven node for each wave.
        self.inputE    =     np.zeros((self.W), dtype=int)   # Index of the driven node for each wave.

        # q  is a (W,N) matrix of generalized wave coordinates (positions)
        # qd is a (W,N) matrix of generalized wave derivatives (velocities)
        self.q  = np.zeros((self.W, self.N))
        self.qd = np.zeros((self.W, self.N))

        # Threshold flags are used to detect events.
        self.rising = np.zeros((self.W, self.N), dtype=bool)
        self.rising_edge_listener = None

        # elapsed integrator time
        self.t = 0.0
        return

[docs]    def set_all_damping(self, b):
        self.b = b*np.ones((self.W))

[docs]    def set_all_period(self, period):
        self.period = period*np.ones((self.W))

[docs]    def get_positions(self):
        """Return a (W, N) matrix of node position values where W is the number of waves and N is the number of samples."""
        return self.q.copy()

[docs]    def update_for_interval(self, interval):
        """Run the simulator for the given interval, which may include one or more integration steps."""
        while interval > 0.0:
            dt = min(interval, 0.001)
            interval -= dt
            self._step(dt)

[docs]    def connect_rising_edge_listener(self, callback):
        self.rising_edge_listener = callback

    # ================================================================
    def _step(self, dt):
        #  calculate the excitation
        if self.t < 0:
            self.qE = np.zeros((self.W))
        else:
            # phase ramps from [0,1) in a sawtooth
            phase = np.mod(self.t, self.period) / self.period

            # create a pulse train with the specified duty cycle
            self.qE = (phase < self.duty) * self.magnitude

        # create an accumulator for forces on all nodes
        force = np.ndarray((self.W, self.N))

        # calculate coupling spring forces on all interior nodes
        #   k (W, 1) * dQ (W, N-2)
        force[:,1:-1] = self.k.reshape((self.W, 1)) * (0.5*(self.q[:, 2:] + self.q[:, 0:-2]) - self.q[:,1:-1])

        # Calculate coupling spring forces on boundary nodes.  This assumes an open boundary condition at each end.
        force[:, 0] = self.k * (self.q[:,  1] - self.q[:, 0])
        force[:,-1] = self.k * (self.q[:, -2] - self.q[:,-1])

        # calculate damping on all nodes
        force -= self.b.reshape((self.W, 1)) * self.qd

        # apply the drive condition as an extra node coupling spring and damper
        force[:, self.inputE] += 4*self.k * (self.qE - self.q[:, self.inputE]) - 20*self.b * (self.qd[:,self.inputE])

        # detect threshold events
        rising = self.q > (0.5 * self.qE.reshape((self.W, 1)))
        new_rising = np.logical_and(rising, np.logical_not(self.rising))
        self.rising = rising
        if self.rising_edge_listener is not None:
            # iterate over the set of indicies of true elements
            for wave, particle in zip(*np.nonzero(new_rising)):
                self.rising_edge_listener(wave, particle)

        # calculate the node dynamics with forward Euler integration
        self.q    += self.qd * dt
        self.qd   += force * dt
        self.t    += dt

        return

################################################################
if __name__ == "__main__":
    obj = Waves1D()
    obj.update_for_interval(0.040)