@startuml
'https://plantuml.com/class-diagram
!theme reddress-lightblue
skinparam groupInheritance 2
skinparam classFontSize 14
skinparam classFontName Arial
skinparam classBackgroundColor %lighten("WhiteSmoke", 5)
skinparam classHeaderBackgroundColor WhiteSmoke
skinparam nodesep 40
skinparam ranksep 12
skinparam tabSize 1

package Light <<Frame>> {
    class Polarization {
        +<color:MediumPurple>ndarray</color> vec
        +<color:MediumPurple>ndarray</color> q_axis
        +<color:MediumPurple>bool</color> vec_as_q
        +<color:MediumPurple>None</color> def_q_axis(<color:MediumPurple>ndarray</color> q_axis, <color:MediumPurple>bool</color> q_fixed)
    }

    class Laser {
        +<color:MediumPurple>float</color> freq
        +<color:MediumPurple>float</color> intensity
        +<color:MediumPurple>Polarization</color> polarization
        +<color:MediumPurple>ndarray</color> k
    }

    Polarization ..> Laser
}

class Environment {
    +<color:MediumPurple>ndarray</color> E
    +<color:MediumPurple>ndarray</color> B
}

package Matter <<Frame>> {

    class State {
        +<color:MediumPurple>float</color> freq_j
        +<color:MediumPurple>float</color> s, l, j, i, f, m
        +<color:MediumPurple>list[float]</color> hyper_const
        +<color:MediumPurple>float</color> g
        +<color:MediumPurple>str</color> label
        +<color:MediumPurple>None</color> update(<color:MediumPurple>Environment</color> environment)
    }

    class DecayMap {
        +<color:MediumPurple>list[tuple[str]]</color> labels
        +<color:MediumPurple>list[float]</color> a
        +<color:MediumPurple>float</color> get(<color:MediumPurple>str</color> label_0, label_1)
    }

    class Atom {
        +<color:MediumPurple>list[State]</color> states
        +<color:MediumPurple>DecayMap</color> decay_map
        +<color:MediumPurple>float</color> mass
        +<color:MediumPurple>float</color> get_state_indexes(<color:MediumPurple>list[str]</color> labels, <color:MediumPurple>list[float]</color> f)
        +<color:MediumPurple>ndarray</color> scattering_rate(<color:MediumPurple>ndarray</color> rho, theta, phi, i, j)
        +<color:MediumPurple>None</color> plot()
    }

    Environment ..> State
    State ..> Atom
    DecayMap ..> Atom
}

class Interaction {
    +<color:MediumPurple>Atom</color> atom
    +<color:MediumPurple>list[Laser]</color> lasers
    +<color:MediumPurple>float</color> delta_max
    +<color:MediumPurple>bool</color> controlled
    +<color:MediumPurple>float</color> dt
    +<color:MediumPurple>None</color> resonance_info()
    +<color:MediumPurple>ndarray</color> rates(<color:MediumPurple>ndarray</color> t, delta v, y0)
    +<color:MediumPurple>ndarray</color> schroedinger(<color:MediumPurple>ndarray</color> t, delta v, y0)
    +<color:MediumPurple>ndarray</color> master(<color:MediumPurple>ndarray</color> t, delta v, y0)
    +<color:MediumPurple>ndarray</color> mc_master(<color:MediumPurple>ndarray</color> t, delta v, y0,
    \t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t<color:MediumPurple>bool</color> dynamics, <color:MediumPurple>int</color> ntraj)
}

Atom .> Interaction
Laser ..> Interaction

@enduml
