From 31e8dde8ce1ec3d9da3b31b5a7afedc13de4d096 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Tue, 12 Nov 2019 00:57:26 -0800 Subject: [PATCH] Add type hints --- meanas/fdtd/energy.py | 59 ++++++++++++++++++++++++++++++++++++------- 1 file changed, 50 insertions(+), 9 deletions(-) diff --git a/meanas/fdtd/energy.py b/meanas/fdtd/energy.py index 7dbf363..9f5c99b 100644 --- a/meanas/fdtd/energy.py +++ b/meanas/fdtd/energy.py @@ -4,7 +4,10 @@ import numpy from .. import dx_lists_t, field_t, field_updater -def poynting(e, h, dxes=None): +def poynting(e: field_t, + h: field_t, + dxes: dx_lists_t = None, + ) -> field_t: if dxes is None: dxes = tuple(tuple(numpy.ones(1) for _ in range(3)) for _ in range(2)) @@ -22,27 +25,52 @@ def poynting(e, h, dxes=None): return s -def poynting_divergence(s=None, *, e=None, h=None, dxes=None): +def poynting_divergence(s: field_t = None, + *, + e: field_t = None, + h: field_t = None, + dxes: dx_lists_t = None, + ) -> field_t: if s is None: s = poynting(e, h, dxes=dxes) ds = ((s[0] - numpy.roll(s[0], 1, axis=0)) + (s[1] - numpy.roll(s[1], 1, axis=1)) + - (s[2] - numpy.roll(s[2], 1, axis=2)) ) + (s[2] - numpy.roll(s[2], 1, axis=2))) return ds -def energy_hstep(e0, h1, e2, epsilon=None, mu=None, dxes=None): +def energy_hstep(e0: field_t, + h1: field_t, + e2: field_t, + epsilon: field_t = None, + mu: field_t = None, + dxes: dx_lists_t = None, + ) -> field_t: u = dxmul(e0 * e2, h1 * h1, epsilon, mu, dxes) return u -def energy_estep(h0, e1, h2, epsilon=None, mu=None, dxes=None): +def energy_estep(h0: field_t, + e1: field_t, + h2: field_t, + epsilon: field_t = None, + mu: field_t = None, + dxes: dx_lists_t = None, + ) -> field_t: u = dxmul(e1 * e1, h0 * h2, epsilon, mu, dxes) return u -def delta_energy_h2e(dt, e0, h1, e2, h3, epsilon=None, mu=None, dxes=None): +def delta_energy_h2e(dt: float, + e0: field_t, + h1: field_t, + e2: field_t, + h3: field_t, + epsilon: field_t = None, + mu: field_t = None, + dxes: dx_lists_t = None, + ) -> field_t: """ This is just from (e2 * e2 + h3 * h1) - (h1 * h1 + e0 * e2) """ @@ -52,7 +80,15 @@ def delta_energy_h2e(dt, e0, h1, e2, h3, epsilon=None, mu=None, dxes=None): return du -def delta_energy_e2h(dt, h0, e1, h2, e3, epsilon=None, mu=None, dxes=None): +def delta_energy_e2h(dt: float, + h0: field_t, + e1: field_t, + h2: field_t, + e3: field_t, + epsilon: field_t = None, + mu: field_t = None, + dxes: dx_lists_t = None, + ) -> field_t: """ This is just from (h2 * h2 + e3 * e1) - (e1 * e1 + h0 * h2) """ @@ -62,7 +98,7 @@ def delta_energy_e2h(dt, h0, e1, h2, e3, epsilon=None, mu=None, dxes=None): return du -def delta_energy_j(j0, e1, dxes=None): +def delta_energy_j(j0: field_t, e1: field_t, dxes: dx_lists_t = None) -> field_t: if dxes is None: dxes = tuple(tuple(numpy.ones(1) for _ in range(3)) for _ in range(2)) @@ -73,7 +109,12 @@ def delta_energy_j(j0, e1, dxes=None): return du -def dxmul(ee, hh, epsilon=None, mu=None, dxes=None): +def dxmul(ee: field_t, + hh: field_t, + epsilon: field_t = None, + mu: field_t = None, + dxes: dx_lists_t = None + ) -> field_t: if epsilon is None: epsilon = 1 if mu is None: