civilquake.py

import tkinter
from matplotlib.figure import Figure
from matplotlib.backend_bases import key_press_handler
from matplotlib.backends.backend_tkagg import (
    FigureCanvasTkAgg,
    NavigationToolbar2Tk
)
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.lines import Line2D
from typing import List, Callable, Tuple


def raising_load(t: np.ndarray, Fmax: float = 1, RT: float = 1) -> callable:
    F = np.zeros_like(t) + Fmax
    t0 = t[0]
    if RT != 0:
        F[t-t0 < RT] = Fmax/RT*(t[t-t0 < RT]-t0)
    return F


raising_load_desc = {
    "Force [N]": "Fmax",
    "Raising time [s]": "RT"
}


class Load:

    def __init__(self,
                 func: Callable,
                 params: dict[str: float],
                 paramsdesc: dict[str: str] = dict()) -> None:
        self.func = func
        self.params = params
        self.paramsdesc = paramsdesc
        if set(self.paramsdesc.values()) != set(self.params.keys()):
            self.paramsdesc = {k: k for k in params}

    def __call__(
                self, t: np.ndarray,
                **params: dict[str: float]
                ) -> np.ndarray:
        if not params:
            params = self.params
        return self.func(t, **params)


class BasicSimulation:

    def __init__(
            self,
            fn: float = 0.5,
            t0: float = 1.0,
            tf: float = 10,
            dt: float = 0.2,
            z: float = 0.01,
            paramsdesc={
                "Natural frequency [Hz]": "fn",
                "Initial time [s]": "t0",
                "Final time [s]": "tf",
                "dt [s]": "dt",
                "Damping [-]": "z"
            },
            load=Load(
                func=raising_load,
                params=dict(Fmax=1, RT=2),
                paramsdesc=raising_load_desc
            ), **kwargs) -> None:

        self.fn = fn
        self.t0 = t0
        self.tf = tf
        self.dt = dt
        self.z = z
        self.paramsdesc = paramsdesc
        self.load = load

        for attribute, value in kwargs.items():
            setattr(self, attribute, value)
        self.recalc()

    def updateAttributes(self) -> np.ndarray:
        self.wn = 2*np.pi*self.fn
        self.Tn = 1/self.fn
        self.n = int((self.tf - self.t0)/self.dt)
        self.t = np.linspace(self.t0, self.tf, self.n)
        self.F = self.load(self.t, **self.load.params)
        return self.F

    def check_stability(self) -> None:
        pass

    def calculateCoefficients(self) -> None:
        raise Warning('Required method not yet implemented')

    def calculate(self) -> None:
        raise Warning('Required method not yet implemented')

    def recalc(self, **kwargs) -> tuple[np.ndarray]:
        for k, v in kwargs.items():
            setattr(self, k, v)
        self.updateAttributes()
        self.check_stability()
        self.calculateCoefficients()
        self.calculate()
        return self.t, self.F, self.x

    def plot(self, show=False, *args, **kwargs) -> List[Line2D]:
        disp = plt.plot(self.t, self.x, self.ls, *args,
                        lw=1, label=self.label, **kwargs)
        force = plt.plot(self.t, self.F, '--r', lw=0.8)
        if show:
            plt.show()
        return disp, force


class Interpolation(BasicSimulation):

    label = "Interpolation"
    ls = "-"

    def calculateCoefficients(self) -> None:
        e = np.exp(-self.z*self.wn*self.dt)
        self.wd = self.wn*np.sqrt(1-self.z**2)
        wdt = self.wd*self.dt
        Rz = np.sqrt(1-self.z**2)

        self.A = e*(self.z/Rz*np.sin(wdt)+np.cos(wdt))
        self.B = e/self.wd*np.sin(wdt)
        self.C = 1/self.wn**2*(
            2*self.z/self.wn/self.dt+e*(
                ((1-2*self.z**2)/wdt-self.z/Rz)*np.sin(wdt) -
                (1+2*self.z/self.wn/self.dt)*np.cos(wdt)
            )
        )
        self.D = 1/self.wn**2*(
            1-2*self.z/(self.wn*self.dt)+e*(
                (2*self.z**2-1)/wdt*np.sin(wdt) +
                2*self.z/self.wn/self.dt*np.cos(wdt)
            )
        )
        self.Ap = -e*(self.wn/Rz*np.sin(wdt))
        self.Bp = e*(np.cos(wdt)-self.z/Rz*np.sin(wdt))
        self.Cp = 1/self.wn**2*(-1/self.dt+e*(
            (self.wn/Rz+self.z/self.dt/Rz)*np.sin(wdt)+1/self.dt*np.cos(wdt)
        ))
        self.Dp = 1/self.wn**2/self.dt*(
            1-e*(self.z/Rz*np.sin(wdt)+np.cos(wdt))
        )

    def calculate(self) -> None:
        u = np.zeros(self.n)
        v = np.zeros(self.n)
        for i, (F, Fnext) in enumerate(zip(self.F[:-1], self.F[1:])):
            u[i+1] = self.A*u[i]+self.B*v[i]+self.C*F+self.D*Fnext
            v[i+1] = self.Ap*u[i]+self.Bp*v[i]+self.Cp*F+self.Dp*Fnext
        self.x = u*self.wn**2


class MCD(BasicSimulation):

    label = "MCD"
    ls = '-.'

    u0: float = 0.0
    v0: float = 0.0

    def check_stability(self) -> None:
        Tn = self.Tn
        self.dt = self.dt
        if Tn/self.dt <= np.pi:
            raise ValueError(f'{Tn/self.dt = } <= pi')

    def calculateCoefficients(self) -> None:
        self.a = np.zeros(self.n)
        self.u = np.zeros_like(self.a)
        self.v = np.zeros_like(self.a)
        self.u[0], self.v[0] = self.u0, self.v0
        self.a[0] = self.F[0]-2*self.z*self.wn*self.v[0]-self.wn**2*self.u[0]
        self.uzero = self.u[0]-self.dt*self.v[0]+self.dt**2/2*self.a[0]
        self.kksurM = 1/self.dt**2+self.z*self.wn/self.dt
        self.asurM = 1/self.dt**2-self.z*self.wn/self.dt
        self.bsurM = self.wn**2-2/self.dt**2

    def calculate(self) -> None:
        self.ppsurM = self.F[0]-self.asurM*self.uzero-self.bsurM*self.u[0]
        self.u[1] = self.ppsurM/self.kksurM
        self.v[0] = self.u[1]-self.uzero/2/self.dt
        self.a[0] = (self.u[1]-2*self.u[0]+self.uzero)/self.dt**2
        for i, Fi in enumerate(self.F[1:-1]):
            i = i+1
            ppsurM = Fi-self.asurM*self.u[i-1]-self.bsurM*self.u[i]
            self.u[i+1] = ppsurM/self.kksurM
            self.v[i] = (self.u[i+1]-self.u[i-1])/2/self.dt
            self.a[i] = (self.u[i+1]-2*self.u[i]+self.u[i-1])/self.dt**2

        self.v[-1] = self.v[-2]
        self.a[-1] = self.a[-2]

        self.x = self.u*self.wn**2


class Newmark(BasicSimulation):

    label = "Newmark"
    ls = '--'

    gamma: float = 1/2
    beta: float = 1/6

    def check_stability(self) -> None:
        Tn, self.dt = self.Tn, self.dt
        gamma, beta = self.gamma, self.beta
        pi, sqrt = np.pi, np.sqrt
        if Tn/self.dt <= np.pi*np.sqrt(2*(gamma-2*beta)):
            raise ValueError(
                f'{Tn/self.dt = } <= {pi*sqrt(2*(gamma-2*beta)) = }'
            )

    def calculateCoefficients(self) -> None:
        self.u = np.zeros(self.n)
        self.v = np.zeros_like(self.u)
        self.a = np.zeros_like(self.u)
        self.a[0] = (
            self.F[0] -
            2*self.z*self.wn*self.v[0] -
            self.wn**2*self.u[0]
        )
        self.kksurM = (
            self.wn**2+1/(self.beta*self.dt**2) +
            2*self.z*self.wn*self.gamma/self.beta/self.dt
        )
        self.asurM = (
            1/self.beta/self.dt +
            2*self.z*self.wn*self.gamma/self.beta
        )
        self.bsurM = 1/2/self.beta+2*self.z*self.wn*self.dt*(
            self.gamma/2/self.beta-1
        )

    def calculate(self) -> None:
        for i, (F, Fnext) in enumerate(zip(self.F[:-1], self.F[1:])):
            dpsurM = (
                Fnext-F +
                self.asurM*self.v[i]+self.bsurM*self.a[i]
            )
            duI = dpsurM/self.kksurM
            dvI = (
                self.gamma/self.beta/self.dt*duI -
                self.gamma/self.beta*self.v[i] +
                self.dt*(1-self.gamma/2/self.beta)*self.a[i]
            )
            daI = (
                1/self.beta/self.dt**2*duI -
                1/self.beta/self.dt*self.v[i] -
                1/2/self.beta*self.a[i]
            )
            self.u[i+1] = self.u[i]+duI
            self.v[i+1] = self.v[i]+dvI
            self.a[i+1] = self.a[i]+daI

            self.x = self.u*self.wn**2


class App:

    def __init__(self,
                 sim: BasicSimulation = Interpolation(),
                 root=None,
                 param_factor=2,
                 param_margin=1,
                 plot_style='bmh'):

        if plot_style:
            plt.style.use(plot_style)
        if root is None:
            root = tkinter.Tk()
        width, height = root.winfo_screenwidth(), root.winfo_screenheight()
        root.geometry(f"{width}x{height}")
        root.title(sim.label)
        self.root = root
        self.sim = sim
        self.pf = param_factor
        self.pm = param_margin

        # Initialize the window, its canvas (figure) and its widgets
        tkinter.Button(
            master=self.root,
            text="Interpolation",
            command=lambda: App(sim=Interpolation(load=self.sim.load))
        ).grid(row=3, column=2)
        self.method = tkinter.Button(
            master=self.root,
            text="Newmark",
            command=lambda: App(sim=Newmark(load=self.sim.load))
        ).grid(row=3, column=3)
        self.method = tkinter.Button(
            master=self.root,
            text="MCD",
            command=lambda: App(sim=MCD(load=self.sim.load))
        ).grid(row=3, column=4)
        self.fig = Figure(figsize=(10, 5))
        self.ax = self.fig.add_subplot(111)
        self.ax.set_xlabel('Time [s]')
        self.ax.set_ylabel('Force [N]\nDisplacement [m]')
        self.loadline, = self.ax.plot(sim.t, sim.F, '-r', mfc='none')
        self.displine, = self.ax.plot(sim.t, sim.x, '-.b')
        self.canvas = FigureCanvasTkAgg(self.fig, self.root)
        toolbar = NavigationToolbar2Tk(
            self.canvas,
            self.root,
            pack_toolbar=False
        )
        toolbar.update()
        self.canvas.mpl_connect("key_press_event", key_press_handler)

        self.canvas.get_tk_widget().grid(row=0, column=0, columnspan=100)
        self.text = tkinter.Text(
            master=self.root,

        )
        toolbar.grid(row=1, column=0)
        self.create_sliders()
        self.create_text()
        self.canvas.draw()
        self.root.mainloop()

    def create_sliders(self):
        params = self.sim.paramsdesc | self.sim.load.paramsdesc
        max_len = max(len(p) for p in params.keys())
        self.sliders = [None for _ in params]

        for i, (k, v) in enumerate(params.items()):
            if v in self.sim.__dict__:
                value = getattr(self.sim, v)
            elif v in self.sim.load.params:
                value = self.sim.load.params[v]
            if value == 0:
                top_value = value + self.pm
            else:
                top_value = self.pf*value

            self.sliders[i] = tkinter.Scale(
                master=self.root,
                from_=0,
                to=top_value,
                resolution=value/100,
                orient=tkinter.HORIZONTAL,
                command=lambda arg, kw=v: self.update((kw, arg)),
                label=f"{k:^{max_len}}"
            )
            self.sliders[i].set(value)
            self.sliders[i].grid(row=1, column=i+1)

    def create_text(self):

        params = self.sim.paramsdesc | self.sim.load.paramsdesc
        self.boxes = [None for _ in params]

        for i, (k, v) in enumerate(params.items()):
            if v in self.sim.__dict__:
                value = getattr(self.sim, v)
            elif v in self.sim.load.params:
                value = self.sim.load.params[v]
            if value == 0:
                top_value = value + 1
            else:
                top_value = 3*value

            self.boxes[i] = tkinter.Entry(
                master=self.root,
                validatecommand=lambda arg: arg.isnumeric(),
                justify="center"
            )
            self.boxes[i].insert(tkinter.END, value)

        butt = tkinter.Button(
            self.root,
            text="Set values",
            command=lambda: self.update(*self.get_text(params))
        )
        butt.grid(row=2, column=0)

        for i, box in enumerate(self.boxes):
            box.grid(row=2, column=i+1)

    def get_text(self, params) -> List[tuple[str]]:
        pairs = [
            (v, box.get())
            for box, (k, v) in zip(self.boxes, params.items())
        ]
        for i, (kw, v) in enumerate(pairs):
            v = float(v)
            self.sliders[i].configure(
                from_=min(0, v),
                to=self.pf*v
            )
            self.sliders[i].set(v)
        return pairs

    def update(self, *args: tuple[str, float]):
        for k, v in args:
            v = float(v)
            if k in self.sim.__dict__:
                setattr(self.sim, k, v)
            elif k in self.sim.load.params:
                self.sim.load.params[k] = v

        for i, (box, slider) in enumerate(zip(self.boxes, self.sliders)):
            self.boxes[i].delete(0, tkinter.END)
            self.boxes[i].insert(tkinter.END, slider.get())

        t, F, x = self.sim.recalc()
        self.loadline.set_data(t, F)
        self.displine.set_data(t, x)
        self.ax.relim()
        self.ax.autoscale_view()
        self.canvas.draw()


if __name__ == '__main__':

    root = tkinter.Tk()
    root.title("My app")

    interpol = Interpolation()
    newmark = Newmark()
    mcd = MCD()

    App(root=root, sim=interpol)