m-thesis-introduction/simulations/lib/location.py

137 lines
3.1 KiB
Python

"""
Locations are wrappers around a Numpy N-dimensional
array.
"""
import numpy as np
from functools import partial
try:
from travelsignal import TravelSignal
except ModuleNotFoundError:
from .travelsignal import TravelSignal
class Location:
"""
A location is a point designated by a spatial coordinate x.
"""
def __init__(self, x):
self.x = np.asarray(x)
def __repr__(self):
return "Location({})".format(repr(self.x))
def __getitem__(self, key):
return self.x[key]
def __setitem__(self, key, val):
self.x[key] = val
# math
def __add__(self, other):
if isinstance(other, Location):
other = other.x
return self.__class__(self.x + other)
def __sub__(self, other):
if isinstance(other, Location):
other = other.x
return self.__class__(self.x - other)
def __mul__(self, other):
return self.__class__(self.x * other)
def __eq__(self, other):
if isinstance(other, Location):
other = other.x
return np.all(self.x == other)
# math alias functions
__radd__ = __add__
__rsub__ = __sub__
__rmul__ = __mul__
class Receiver(Location):
"""
A location able to trace a signal over time.
Optionally applies a transformation to the traced signal.
"""
def __repr__(self):
return "Receiver({})".format(repr(self.x))
def recv(self, travel_signal: TravelSignal) -> TravelSignal:
"""
Return a function that traces the signal as a function of time
at the receiver's location
"""
return partial(travel_signal, x_f=self.x)
receive = recv
class Emitter(Location):
"""
Emit a signal from position x_0 (and time t_0)
"""
def __repr__(self):
return "Emitter({})".format(repr(self.x))
def emit(self, travel_signal: TravelSignal) -> TravelSignal:
return partial(travel_signal, x_0=self.x)
if __name__ == "__main__":
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
# 2D showcase
source = Emitter([1,1])
antennae = [
Receiver([2,3]),
Receiver([10,10]),
Receiver([-2,-3]),
Receiver([-10,0]),
]
fig, ax = plt.subplots()
ax.set_title("Geometry of Emitter(s) and Antennae")
ax.set_ylabel("y")
ax.set_xlabel("x")
ax.plot(*source.x, '*', label="Emitter")
for j, ant in enumerate(antennae):
ax.plot(*ant.x, '+', label="Antenna {}".format(j))
ax.legend()
fig.show()
# 3D showcase
source = Emitter([1,1,1])
antennae = [
Receiver([2,3,0]),
Receiver([10,10,-5]),
Receiver([-2,-3,9]),
Receiver([-10,0,-5]),
]
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_title("Geometry of Emitter(s) and Antennae")
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
ax.plot([source.x[0]], *source.x[1:], '*', label="Emitter")
for j, ant in enumerate(antennae):
ax.plot([ant.x[0]], *ant.x[1:], '+', label="Antenna {}".format(j))
ax.legend()
plt.show()