def __init__(self, params):
self.init = False
self.connected_to_atmosphere = False
self.params = params
natm, noc = self.params.nmod
self.M = sp.zeros((noc, noc), dtype=float, format='dok')
self.N = sp.zeros((noc, noc), dtype=float, format='dok')
self.O = sp.zeros((noc, noc, noc), dtype=float, format='dok')
self.C = sp.zeros((noc, noc, noc), dtype=float, format='dok')
self.K = None
self.W = None
# initialization of the variables
oceanic_wavenumbers = np.empty(noc, dtype=object)
# Oceanic wavenumbers definition
oms = self.params.goblocks
for i in range(
oms.shape[0]
): # function type is limited to L for the moment: ocean is a closed basin
oceanic_wavenumbers[i] = WaveNumber('L', oms[i, 1], 0, oms[i, 0],
oms[i, 0] / 2., oms[i, 1])
self.oceanic_wavenumbers = oceanic_wavenumbers
self._calculate_M()
self._calculate_N()
self._calculate_O()
self._calculate_C()
def connect_to_atmosphere(self, atmosphere_inner_products):
"""Connect the ocean to an atmosphere.
Parameters
----------
atmosphere_inner_products: AtmosphericAnalyticInnerProducts
The inner products of the atmosphere.
"""
self.atmosphere_inner_products = atmosphere_inner_products
self.connected_to_atmosphere = True
if self.stored:
natm = atmosphere_inner_products.natm
self._K = sp.zeros((self.noc, natm), dtype=float, format='dok')
self._W = sp.zeros((self.noc, natm), dtype=float, format='dok')
args_list = [(i, j) for i in range(self.noc) for j in range(natm)]
for arg in args_list:
# K inner products
self._K[arg] = self._K_comp(*arg)
# W inner products
self._W[arg] = self._W_comp(*arg)
self._K = self._K.to_coo()
self._W = self._W.to_coo()
self.atmosphere_inner_products = None
def connect_to_ocean(self, ocean_basis, num_threads=None):
"""Connect the atmosphere to an ocean.
Parameters
----------
ocean_basis: SymbolicBasis or OceanicSymbolicInnerProducts
Basis of function of the ocean or a symbolic oceanic inner products object containing the basis.
num_threads: int or None, optional
Number of threads to use to compute the symbolic inner products. If `None` use all the cpus available.
Default to `None`.
"""
if isinstance(ocean_basis, OceanicSymbolicInnerProducts):
ocean_basis = ocean_basis.oceanic_basis
self.ground_basis = None
self.connected_to_ground = False
self.oceanic_basis = ocean_basis
self.connected_to_ocean = True
if self.stored:
if num_threads is None:
pool = multiprocessing.Pool(
processes=multiprocessing.cpu_count())
else:
pool = multiprocessing.Pool(processes=num_threads)
noc = len(ocean_basis)
self._gh = None
self._d = sp.zeros((self.natm, noc), dtype=float, format='dok')
self._s = sp.zeros((self.natm, noc), dtype=float, format='dok')
# d inner products
args_list = [[(i, j), self.iip.ip_lap, (self._F(i), self._phi(j))]
for i in range(self.natm) for j in range(noc)]
result = pool.map(_apply, args_list)
for res in result:
self._d[res[0]] = float(res[1].subs(self.subs).subs(
self.atmospheric_basis.substitutions).subs(
self.oceanic_basis.substitutions))
# s inner products
args_list = [[(i, j), self.iip.symbolic_inner_product,
(self._F(i), self._phi(j))] for i in range(self.natm)
for j in range(noc)]
result = pool.map(_apply, args_list)
for res in result:
self._s[res[0]] = float(res[1].subs(self.subs).subs(
self.atmospheric_basis.substitutions).subs(
self.oceanic_basis.substitutions))
self._s = self._s.to_coo()
self._d = self._d.to_coo()
pool.terminate()
def __init__(self, params):
self.connected_to_ocean = False
self.params = params
natm, noc = self.params.nmod
if natm == 0:
exit("*** Problem with inner products : natm==0!***")
self.a = sp.zeros((natm, natm), dtype=float, format='dok')
self.c = sp.zeros((natm, natm), dtype=float, format='dok')
self.b = sp.zeros((natm, natm, natm), dtype=float, format='dok')
self.g = sp.zeros((natm, natm, natm), dtype=float, format='dok')
self.d = None
self.s = None
# initialization of the variables
atmospheric_wavenumbers = np.empty(natm, dtype=object)
j = -1
# Atmospheric wavenumbers definition
ams = self.params.ablocks
for i in range(
ams.shape[0]
): # function type is limited to AKL for the moment: atmosphere is a channel
if ams[i, 0] == 1:
atmospheric_wavenumbers[j + 1] = WaveNumber(
'A', ams[i, 1], 0, 0, 0, ams[i, 1])
atmospheric_wavenumbers[j + 2] = WaveNumber(
'K', ams[i, 1], ams[i, 0], 0, ams[i, 0], ams[i, 1])
atmospheric_wavenumbers[j + 3] = WaveNumber(
'L', ams[i, 1], 0, ams[i, 0], ams[i, 0], ams[i, 1])
j = j + 3
else:
atmospheric_wavenumbers[j + 1] = WaveNumber(
'K', ams[i, 1], ams[i, 0], 0, ams[i, 0], ams[i, 1])
atmospheric_wavenumbers[j + 2] = WaveNumber(
'L', ams[i, 1], 0, ams[i, 0], ams[i, 0], ams[i, 1])
j = j + 2
self.atmospheric_wavenumbers = atmospheric_wavenumbers
self._calculate_a()
self._calculate_g()
self._calculate_b()
self._calculate_c()
def test_indexrandom2(self):
M1 = sparse.zeros(shape=10)
M1[:10] = 2
print(M1.to_numpy())
M2 = sparse.zeros(shape=(10, 10))
M2[:, 3] = M1
print(M2.to_numpy())
M2[3, :] = M1
print(M2.to_numpy())
def connect_to_atmosphere(self, atmosphere_basis, num_threads=None):
"""Connect the ocean to an atmosphere.
Parameters
----------
atmosphere_basis: SymbolicBasis or AtmosphericSymbolicInnerProducts
Basis of function of the atmosphere or a symbolic atmospheric inner products object containing the basis.
num_threads: int or None, optional
Number of threads to use to compute the symbolic inner products. If `None` use all the cpus available.
Default to `None`.
"""
if isinstance(atmosphere_basis, AtmosphericSymbolicInnerProducts):
atmosphere_basis = atmosphere_basis.atmospheric_basis
self.atmospheric_basis = atmosphere_basis
self.connected_to_atmosphere = True
if self.stored:
if num_threads is None:
pool = multiprocessing.Pool(
processes=multiprocessing.cpu_count())
else:
pool = multiprocessing.Pool(processes=num_threads)
natm = len(atmosphere_basis)
self._K = sp.zeros((self.noc, natm), dtype=float, format='dok')
self._W = sp.zeros((self.noc, natm), dtype=float, format='dok')
# K inner products
l = [[(i, j), self.iip.ip_lap, (self._phi(i), self._F(j))]
for i in range(self.noc) for j in range(natm)]
result = pool.map(_apply, l)
for res in result:
self._K[res[0]] = float(res[1].subs(self.subs).subs(
self.oceanic_basis.substitutions).subs(
self.atmospheric_basis.substitutions))
# W inner products
l = [[(i, j), self.iip.symbolic_inner_product,
(self._phi(i), self._F(j))] for i in range(self.noc)
for j in range(natm)]
result = pool.map(_apply, l)
for res in result:
self._W[res[0]] = float(res[1].subs(self.subs).subs(
self.oceanic_basis.substitutions).subs(
self.atmospheric_basis.substitutions))
self._K = self._K.to_coo()
self._W = self._W.to_coo()
pool.terminate()
def connect_to_ground(self, ground_inner_products):
"""Connect the atmosphere to the ground.
Parameters
----------
ground_inner_products: GroundAnalyticInnerProducts
The inner products of the ground.
"""
self.ocean_inner_products = None
self.connected_to_ocean = False
self.ground_inner_products = ground_inner_products
self.connected_to_ground = True
if self.stored:
ngr = ground_inner_products.ngr
self._s = sp.zeros((self.natm, ngr), dtype=float, format='dok')
args_list = [(i, j) for i in range(self.natm) for j in range(ngr)]
# s inner products
for arg in args_list:
self._s[arg] = self._s_comp(*arg)
self._s = self._s.to_coo()
# ensure that the ocean is connected as well
if not ground_inner_products.connected_to_atmosphere:
ground_inner_products.connect_to_atmosphere(self)
if self.stored:
self.ground_inner_products = None
def compute_inner_products(self, num_threads=None):
"""Function computing and storing all the inner products at once.
Parameters
----------
num_threads: int or None, optional
Number of threads to use to compute the symbolic inner products. If `None` use all the cpus available.
Default to `None`.
"""
self._U = sp.zeros((self.ngr, self.ngr), dtype=float, format='dok')
if num_threads is None:
pool = multiprocessing.Pool(processes=multiprocessing.cpu_count())
else:
pool = multiprocessing.Pool(processes=num_threads)
# U inner products
l = [[(i, j), self.ip.symbolic_inner_product,
(self._phi(i), self._phi(j))] for i in range(self.ngr)
for j in range(self.ngr)]
result = pool.map(_apply, l)
for res in result:
self._U[res[0]] = float(res[1].subs(self.subs).subs(
self.ground_basis.substitutions))
pool.terminate()
def empty_mask(mask_shape, dtype):
'''
Return an empty sparse mask
Improve readability
'''
return sparse.zeros(mask_shape, dtype=dtype)
def masks():
input_masks = [
lambda: np.ones((128, 128)),
lambda: sparse.zeros((128, 128)),
lambda: np.ones((128, 128)),
lambda: sp.csr_matrix(((1,), ((64,), (64,))), shape=(128, 128), dtype=np.float32),
lambda: gradient_x(128, 128, dtype=np.float32),
]
return MaskContainer(mask_factories=input_masks, dtype=np.float32)
def test_indexrandom1(self):
M1 = sparse.zeros(shape=(10, 10, 10))
M1[:10:2, 5, 0:10:3] = 2
M1[:10:2, 5, 5:10:2] = 3
print(M1.to_numpy())
M2 = np.zeros(shape=(10, 10, 10))
M2[:10:2, 5, 0:10:3] = 2
M2[:10:2, 5, 5:10:2] = 3
print(M2)
self.assertTrue(np.array_equal(M1.to_numpy(), M2))
def compute_inner_products(self):
"""Function computing and storing all the inner products at once."""
self._U = sp.zeros((self.ngr, self.ngr), dtype=float, format='dok')
args_list = [(i, j) for i in range(self.ngr) for j in range(self.ngr)]
for arg in args_list:
# U inner products
self._U[arg] = self._U_comp(*arg)
self._U = self._U.to_coo()
def test_mask_patch_sparse():
for i in range(REPEATS):
print(f"Loop number {i}")
num_nav_dims = np.random.choice([1, 2, 3])
num_sig_dims = np.random.choice([2, 3])
nav_dims = tuple(np.random.randint(low=8, high=16, size=num_nav_dims))
sig_dims = tuple(np.random.randint(low=8, high=16, size=num_sig_dims))
# The mask-based correction is performed as float64 since it creates
# numerical instabilities otherwise
data = gradient_data(nav_dims, sig_dims).astype(np.float64)
gain_map = (np.random.random(sig_dims) + 1).astype(np.float64)
dark_image = np.random.random(sig_dims).astype(np.float64)
exclude = exclude_pixels(sig_dims=sig_dims, num_excluded=3)
damaged_data = data.copy()
damaged_data /= gain_map
damaged_data += dark_image
damaged_data[(Ellipsis, *exclude)] = 1e24
print("Nav dims: ", nav_dims)
print("Sig dims:", sig_dims)
print("Exclude: ", exclude)
masks = sparse.DOK(sparse.zeros((20, ) + sig_dims, dtype=np.float64))
indices = [
np.random.randint(low=0, high=s, size=s // 2)
for s in (20, ) + sig_dims
]
for tup in zip(*indices):
masks[tup] = 1
masks = masks.to_coo()
data_flat = data.reshape((np.prod(nav_dims), np.prod(sig_dims)))
damaged_flat = damaged_data.reshape(
(np.prod(nav_dims), np.prod(sig_dims)))
correct_dot = sparse.dot(data_flat,
masks.reshape((-1, np.prod(sig_dims))).T)
corrected_masks = detector.correct_dot_masks(masks, gain_map, exclude)
assert is_sparse(corrected_masks)
reconstructed_dot =\
sparse.dot(damaged_flat, corrected_masks.reshape((-1, np.prod(sig_dims))).T)\
- sparse.dot(dark_image.flatten(), corrected_masks.reshape((-1, np.prod(sig_dims))).T)
_check_result(data=correct_dot,
corrected=reconstructed_dot,
atol=1e-8,
rtol=1e-5)
def connect_to_ocean(self, ocean_inner_products):
"""Connect the atmosphere to an ocean.
Parameters
----------
ocean_inner_products: OceanicInnerProducts
The inner products of the ocean.
"""
natm, noc = self.params.nmod
self.d = sp.zeros((natm, noc), dtype=float, format='dok')
self.s = sp.zeros((natm, noc), dtype=float, format='dok')
self._calculate_s(ocean_inner_products)
# ensure that the ocean is connected as well
if not ocean_inner_products.connected_to_atmosphere:
ocean_inner_products.connect_to_atmosphere(self)
self._calculate_d(ocean_inner_products)
self.connected_to_ocean = True
def compute_inner_products(self):
"""Function computing and storing all the inner products at once."""
self._a = sp.zeros((self.natm, self.natm), dtype=float, format='dok')
self._u = sp.zeros((self.natm, self.natm), dtype=float, format='dok')
self._c = sp.zeros((self.natm, self.natm), dtype=float, format='dok')
self._b = sp.zeros((self.natm, self.natm, self.natm),
dtype=float,
format='dok')
self._g = sp.zeros((self.natm, self.natm, self.natm),
dtype=float,
format='dok')
args_list = [(i, j) for i in range(self.natm)
for j in range(self.natm)]
for arg in args_list:
# a inner products
self._a[arg] = self._a_comp(*arg)
# u inner products
self._u[arg] = self._u_comp(*arg)
# c inner products
self._c[arg] = self._c_comp(*arg)
args_list = [(i, j, k) for i in range(self.natm)
for j in range(self.natm) for k in range(self.natm)]
for arg in args_list:
# g inner products
val = self._g_comp(*arg)
self._g[arg] = val
# b inner products
self._b[arg] = val * self._a[arg[-1], arg[-1]]
self._a = self._a.to_coo()
self._u = self._u.to_coo()
self._c = self._c.to_coo()
self._g = self._g.to_coo()
self._b = self._b.to_coo()
def compute_inner_products(self):
"""Function computing and storing all the inner products at once."""
self._M = sp.zeros((self.noc, self.noc), dtype=float, format='dok')
self._U = sp.zeros((self.noc, self.noc), dtype=float, format='dok')
self._N = sp.zeros((self.noc, self.noc), dtype=float, format='dok')
self._O = sp.zeros((self.noc, self.noc, self.noc),
dtype=float,
format='dok')
self._C = sp.zeros((self.noc, self.noc, self.noc),
dtype=float,
format='dok')
args_list = [(i, j) for i in range(self.noc) for j in range(self.noc)]
for arg in args_list:
# M inner products
self._M[arg] = self._M_comp(*arg)
# U inner products
self._U[arg] = self._U_comp(*arg)
# N inner products
self._N[arg] = self._N_comp(*arg)
args_list = [(i, j, k) for i in range(self.noc)
for j in range(self.noc) for k in range(self.noc)]
for arg in args_list:
# O inner products
val = self._O_comp(*arg)
self._O[arg] = val
# C inner products
self._C[arg] = val * self._M[arg[-1], arg[-1]]
self._M = self._M.to_coo()
self._U = self._U.to_coo()
self._N = self._N.to_coo()
self._O = self._O.to_coo()
self._C = self._C.to_coo()
def connect_to_atmosphere(self, atmosphere_inner_products):
"""Connect the ocean to an atmosphere.
Parameters
----------
atmosphere_inner_products: AtmosphericInnerProducts
The inner products of the atmosphere.
"""
natm, noc = self.params.nmod
self.K = sp.zeros((noc, natm), dtype=float, format='dok')
self.W = sp.zeros((noc, natm), dtype=float, format='dok')
if atmosphere_inner_products.s is None:
atmosphere_inner_products.s = sp.zeros((natm, noc),
dtype=float,
format='dok')
atmosphere_inner_products._calculate_s(self)
self._calculate_W(atmosphere_inner_products)
self._calculate_K(atmosphere_inner_products)
self.connected_to_atmosphere = True
def connect_to_ocean(self, ocean_inner_products):
"""Connect the atmosphere to an ocean.
Parameters
----------
ocean_inner_products: OceanicInnerAnalyticProducts
The inner products of the ocean.
"""
self.ground_inner_products = None
self.connected_to_ground = False
self.ocean_inner_products = ocean_inner_products
self.connected_to_ocean = True
if self.stored:
noc = ocean_inner_products.noc
self._s = sp.zeros((self.natm, noc), dtype=float, format='dok')
self._d = sp.zeros((self.natm, noc), dtype=float, format='dok')
args_list = [(i, j) for i in range(self.natm) for j in range(noc)]
for arg in args_list:
# s inner products
self._s[arg] = self._s_comp(*arg)
# d inner products
self._d[arg] = self._d_comp(*arg)
self._s = self._s.to_coo()
self._d = self._d.to_coo()
# ensure that the ocean is connected as well
if not ocean_inner_products.connected_to_atmosphere:
ocean_inner_products.connect_to_atmosphere(self)
if self.stored:
self.ocean_inner_products = None
def test_zeros(self):
M1 = sparse.zeros(shape=(30, 30))
M2 = np.zeros((30, 30))
self.assertTrue(np.array_equal(M1.to_numpy(), M2))
self.assertEqual(M1.T.shape[0], 0)
def compute_tensor(self):
"""Routine to compute the tensor."""
if self.params is None:
return
if self.atmospheric_inner_products is None and self.oceanic_inner_products is None \
and self.ground_inner_products is None:
return
aips = self.atmospheric_inner_products
par = self.params
atp = par.atemperature_params
ap = par.atmospheric_params
op = par.oceanic_params
scp = par.scale_params
gp = par.ground_params
namod = par.nmod[0]
ngomod = par.nmod[1]
ndim = par.ndim
bips = None
if self.oceanic_inner_products is not None:
bips = self.oceanic_inner_products
ocean = True
else:
ocean = False
if self.ground_inner_products is not None:
bips = self.ground_inner_products
ground_temp = True
else:
ground_temp = False
# 0-th tensor component is an empty matrix
tensor = sp.zeros((ndim+1, ndim + 1, ndim + 1), dtype=np.float64, format='dok')
jacobian_tensor = sp.zeros((ndim+1, ndim + 1, ndim + 1), dtype=np.float64, format='dok')
# constructing some derived matrices
if aips is not None:
a_inv = np.zeros((namod, namod))
for i in range(namod):
for j in range(namod):
a_inv[i, j] = aips.a(i, j)
a_inv = np.linalg.inv(a_inv)
a_theta = np.zeros((namod, namod))
for i in range(namod):
for j in range(namod):
a_theta[i, j] = ap.sig0 * aips.a(i, j) - aips.u(i, j)
a_theta = np.linalg.inv(a_theta)
if bips is not None:
U_inv = np.zeros((ngomod, ngomod))
for i in range(ngomod):
for j in range(ngomod):
U_inv[i, j] = bips.U(i, j)
U_inv = np.linalg.inv(U_inv)
if ocean:
M_psio = np.zeros((ngomod, ngomod))
for i in range(ngomod):
for j in range(ngomod):
M_psio[i, j] = bips.M(i, j) + par.G * bips.U(i, j)
M_psio = np.linalg.inv(M_psio)
#################
if bips is not None:
go = bips.stored
else:
go = True
if aips.stored and go:
# psi_a part
for i in range(1, namod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
for j in range(1, namod + 1):
val = a_inv[(i - 1), :] @ aips._c[:, (j - 1)]
t[self._psi_a(j), 0] -= val * scp.beta
t[self._psi_a(j), 0] -= (ap.kd * _kronecker_delta((i - 1), (j - 1))) / 2
t[self._theta_a(j), 0] = (ap.kd * _kronecker_delta((i - 1), (j - 1))) / 2
if gp is not None:
if gp.hk is not None:
if gp.orographic_basis == "atmospheric":
oro = a_inv[(i - 1), :] @ aips._g[:, (j - 1), :] @ gp.hk
else:
oro = a_inv[(i - 1), :] @ aips._gh[:, (j - 1), :] @ gp.hk
t[self._psi_a(j), 0] -= oro / 2
t[self._theta_a(j), 0] += oro / 2
for k in range(1, namod + 1):
val = a_inv[(i - 1), :] @ aips._b[:, (j - 1), (k - 1)]
t[self._psi_a(j), self._psi_a(k)] = - val
t[self._theta_a(j), self._theta_a(k)] = - val
if ocean:
for j in range(1, ngomod + 1):
val = a_inv[(i - 1), :] @ aips._d[:, (j - 1)]
t[self._psi_o(j), 0] += val * ap.kd / 2
t = self.simplify_matrix(t)
tensor[self._psi_a(i)] = t
jacobian_tensor[self._psi_a(i)] = t + t.T
# theta_a part
for i in range(1, namod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
if par.Cpa is not None:
t[0, 0] -= a_theta[(i - 1), :] @ aips._u @ par.Cpa
if atp.hd is not None and atp.thetas is not None:
val = - a_theta[(i - 1), :] @ aips._u @ atp.thetas
t[0, 0] += val * atp.hd
for j in range(1, namod + 1):
val = a_theta[(i - 1), :] @ aips._a[:, (j - 1)]
t[self._psi_a(j), 0] += val * ap.kd * ap.sig0 / 2
t[self._theta_a(j), 0] -= val * (ap.kd / 2 + 2 * ap.kdp) * ap.sig0
val = - a_theta[(i - 1), :] @ aips._c[:, (j - 1)]
t[self._theta_a(j), 0] += val * scp.beta * ap.sig0
val = a_theta[(i - 1), :] @ aips._u[:, (j - 1)]
if par.LSBpa is not None and par.Lpa is not None:
t[self._theta_a(j), 0] += val * (par.LSBpa + atp.sc * par.Lpa)
if atp.hd is not None:
t[self._theta_a(j), 0] += val * atp.hd
if gp is not None:
if gp.hk is not None:
if gp.orographic_basis == "atmospheric":
oro = a_theta[(i - 1), :] @ aips._g[:, (j - 1), :] @ gp.hk
else:
oro = a_theta[(i - 1), :] @ aips._gh[:, (j - 1), :] @ gp.hk
t[self._theta_a(j), 0] -= ap.sig0 * oro / 2
t[self._psi_a(j), 0] += ap.sig0 * oro / 2
for k in range(1, namod + 1):
val = a_theta[(i - 1), :] @ aips._b[:, (j - 1), (k - 1)]
t[self._psi_a(j), self._theta_a(k)] = - val * ap.sig0
t[self._theta_a(j), self._psi_a(k)] = - val * ap.sig0
val = a_theta[(i - 1), :] @ aips._g[:, (j - 1), (k - 1)]
t[self._psi_a(j), self._theta_a(k)] += val
if ocean:
for j in range(1, ngomod + 1):
val = - a_theta[(i - 1), :] @ aips._d[:, (j - 1)]
t[self._psi_o(j), 0] += val * ap.sig0 * ap.kd / 2
if par.LSBpgo is not None and par.Lpa is not None:
val = - a_theta[(i - 1), :] @ aips._s[:, (j - 1)]
t[self._deltaT_o(j), 0] += val * (par.LSBpgo + par.Lpa / 2)
if ground_temp:
for j in range(1, ngomod + 1):
if par.LSBpgo is not None and par.Lpa is not None:
val = - a_theta[(i - 1), :] @ aips._s[:, (j - 1)]
t[self._deltaT_g(j), 0] += val * (par.LSBpgo + par.Lpa / 2)
t = self.simplify_matrix(t)
tensor[self._theta_a(i)] = t
jacobian_tensor[self._theta_a(i)] = t + t.T
if ocean:
# psi_o part
for i in range(1, ngomod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
for j in range(1, namod + 1):
val = M_psio[(i - 1), :] @ bips._K[:, (j - 1)] * op.d
t[self._psi_a(j), 0] += val
t[self._theta_a(j), 0] -= val
for j in range(1, ngomod + 1):
val = - M_psio[(i - 1), :] @ bips._N[:, (j - 1)]
t[self._psi_o(j), 0] += val * scp.beta
val = - M_psio[(i - 1), :] @ bips._M[:, (j - 1)]
t[self._psi_o(j), 0] += val * (op.r + op.d)
for k in range(1, ngomod + 1):
t[self._psi_o(j), self._psi_o(k)] -= M_psio[(i - 1), :] @ bips._C[:, (j - 1), (k - 1)]
t = self.simplify_matrix(t)
tensor[self._psi_o(i)] = t
jacobian_tensor[self._psi_o(i)] = t + t.T
# deltaT_o part
for i in range(1, ngomod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
t[0, 0] += bips._W[(i - 1), :] @ np.array(par.Cpgo, dtype=np.float64) # problem with sparse matmul and object dtype
for j in range(1, namod + 1):
val = U_inv[(i - 1), :] @ bips._W[:, (j - 1)]
t[self._theta_a(j), 0] += val * (2 * atp.sc * par.Lpgo + par.sbpa)
for j in range(1, ngomod + 1):
t[self._deltaT_o(j), 0] = - (par.Lpgo + par.sbpgo) * _kronecker_delta((i - 1), (j - 1))
for k in range(1, ngomod + 1):
t[self._psi_o(j), self._deltaT_o(k)] -= U_inv[(i - 1), :] @ bips._O[:, (j - 1), (k - 1)]
t = self.simplify_matrix(t)
tensor[self._deltaT_o(i)] = t
jacobian_tensor[self._deltaT_o(i)] = t + t.T
# deltaT_g part
if ground_temp:
for i in range(1, ngomod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
t[0, 0] += bips._W[(i - 1), :] @ np.array(par.Cpgo, dtype=np.float64)
for j in range(1, namod + 1):
val = U_inv[(i - 1), :] @ bips._W[:, (j - 1)]
t[self._theta_a(j), 0] += val * (2 * atp.sc * par.Lpgo + par.sbpa)
for j in range(1, ngomod + 1):
t[self._deltaT_g(j), 0] = - (par.Lpgo + par.sbpgo) * _kronecker_delta((i - 1), (j - 1))
t = self.simplify_matrix(t)
tensor[self._deltaT_g(i)] = t
jacobian_tensor[self._deltaT_g(i)] = t + t.T
self.tensor = tensor.to_coo()
self.jacobian_tensor = jacobian_tensor.to_coo()
else:
# psi_a part
for i in range(1, namod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
for j in range(1, namod + 1):
val = 0
for jj in range(1, namod + 1):
val -= a_inv[(i - 1), (jj - 1)] * aips.c((jj - 1), (j - 1))
t[self._psi_a(j), 0] += val * scp.beta
t[self._psi_a(j), 0] -= (ap.kd * _kronecker_delta((i - 1), (j - 1))) / 2
t[self._theta_a(j), 0] = (ap.kd * _kronecker_delta((i - 1), (j - 1))) / 2
if gp is not None:
if gp.hk is not None:
oro = 0
if gp.orographic_basis == "atmospheric":
for jj in range(1, namod + 1):
for k in range(1, namod + 1):
oro += a_inv[(i - 1), (jj - 1)] * aips.g((jj - 1), (j - 1), (k - 1)) * gp.hk[(k - 1)]
else:
for jj in range(1, namod + 1):
for k in range(1, ngomod + 1):
oro += a_inv[(i - 1), (jj - 1)] * aips.gh((jj - 1), (j - 1), (k - 1)) * gp.hk[(k - 1)]
t[self._psi_a(j), 0] -= oro / 2
t[self._theta_a(j), 0] += oro / 2
for k in range(1, namod + 1):
val = 0
for jj in range(1, namod + 1):
val += a_inv[(i - 1), (jj - 1)] * aips.b((jj - 1), (j - 1), (k - 1))
t[self._psi_a(j), self._psi_a(k)] = - val
t[self._theta_a(j), self._theta_a(k)] = - val
if ocean:
for j in range(1, ngomod + 1):
val = 0
for jj in range(1, namod + 1):
val += a_inv[(i - 1), (jj - 1)] * aips.d((jj - 1), (j - 1))
t[self._psi_o(j), 0] += val * ap.kd / 2
t = self.simplify_matrix(t)
tensor[self._psi_a(i)] = t
jacobian_tensor[self._psi_a(i)] = t + t.T
# theta_a part
for i in range(1, namod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
if par.Cpa is not None:
for jj in range(1, namod + 1):
for kk in range(1, namod + 1):
t[0, 0] -= a_theta[(i - 1), (jj - 1)] * aips.u((jj - 1), (kk - 1)) * par.Cpa[kk - 1]
if atp.hd is not None and atp.thetas is not None:
val = 0
for jj in range(1, namod + 1):
for kk in range(1, namod + 1):
val -= a_theta[(i - 1), (jj - 1)] * aips.u((jj - 1), (kk - 1)) * atp.thetas[(kk - 1)]
t[0, 0] += val * atp.hd
for j in range(1, namod + 1):
val = 0
for jj in range(1, namod + 1):
val += a_theta[(i - 1), (jj - 1)] * aips.a((jj - 1), (j - 1))
t[self._psi_a(j), 0] += val * ap.kd * ap.sig0 / 2
val = 0
for jj in range(1, namod + 1):
val -= a_theta[(i - 1), (jj - 1)] * aips.c((jj - 1), (j - 1))
t[self._theta_a(j), 0] += val * scp.beta * ap.sig0
val = 0
for jj in range(1, namod + 1):
val -= a_theta[(i - 1), (jj - 1)] * aips.a((jj - 1), (j - 1))
t[self._theta_a(j), 0] += val * (ap.kd / 2 + 2 * ap.kdp) * ap.sig0
if par.LSBpa is not None and par.Lpa is not None:
val = 0
for jj in range(1, namod + 1):
val += a_theta[(i - 1), (jj - 1)] * aips.u((jj - 1), (j - 1))
t[self._theta_a(j), 0] += val * (par.LSBpa + atp.sc * par.Lpa)
if atp.hd is not None:
val = 0
for jj in range(1, namod + 1):
val += a_theta[(i - 1), (jj - 1)] * aips.u((jj - 1), (j - 1))
t[self._theta_a(j), 0] += val * atp.hd
if gp is not None:
if gp.hk is not None:
oro = 0
if gp.orographic_basis == "atmospheric":
for jj in range(1, namod + 1):
for k in range(1, namod + 1):
oro += a_theta[(i - 1), (jj - 1)] * aips.g((jj - 1), (j - 1), (k - 1)) * gp.hk[(k - 1)]
else:
for jj in range(1, namod + 1):
for k in range(1, ngomod + 1):
oro += a_theta[(i - 1), (jj - 1)] * aips.gh((jj - 1), (j - 1), (k - 1)) * gp.hk[(k - 1)]
t[self._theta_a(j), 0] -= ap.sig0 * oro / 2
t[self._psi_a(j), 0] += ap.sig0 * oro / 2
for k in range(1, namod + 1):
val = 0
for jj in range(1, namod + 1):
val += a_theta[(i - 1), (jj - 1)] * aips.b((jj - 1), (j - 1), (k - 1))
t[self._psi_a(j), self._theta_a(k)] = - val * ap.sig0
t[self._theta_a(j), self._psi_a(k)] = - val * ap.sig0
val = 0
for jj in range(1, namod + 1):
val += a_theta[(i - 1), (jj - 1)] * aips.g((jj - 1), (j - 1), (k - 1))
t[self._psi_a(j), self._theta_a(k)] += val
if ocean:
for j in range(1, ngomod + 1):
val = 0
for jj in range(1, namod + 1):
val -= a_theta[(i - 1), (jj - 1)] * aips.d((jj - 1), (j - 1))
t[self._psi_o(j), 0] += val * ap.sig0 * ap.kd / 2
if par.LSBpgo is not None and par.Lpa is not None:
val = 0
for jj in range(1, namod + 1):
val -= a_theta[(i - 1), (jj - 1)] * aips.s((jj - 1), (j - 1))
t[self._deltaT_o(j), 0] += val * (par.LSBpgo + par.Lpa / 2)
if ground_temp:
for j in range(1, ngomod + 1):
if par.LSBpgo is not None and par.Lpa is not None:
val = 0
for jj in range(1, namod + 1):
val -= a_theta[(i - 1), (jj - 1)] * aips.s((jj - 1), (j - 1))
t[self._deltaT_g(j), 0] += val * (par.LSBpgo + par.Lpa / 2)
t = self.simplify_matrix(t)
tensor[self._theta_a(i)] = t
jacobian_tensor[self._theta_a(i)] = t + t.T
if ocean:
# psi_o part
for i in range(1, ngomod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
for j in range(1, namod + 1):
for jj in range(1, ngomod + 1):
val = M_psio[(i - 1), (jj - 1)] * bips.K((jj - 1), (j - 1)) * op.d
t[self._psi_a(j), 0] += val
t[self._theta_a(j), 0] -= val
for j in range(1, ngomod + 1):
val = 0
for jj in range(1, ngomod + 1):
val -= M_psio[(i - 1), (jj - 1)] * bips.N((i - 1), (j - 1))
t[self._psi_o(j), 0] += val * scp.beta
val = 0
for jj in range(1, ngomod + 1):
val -= M_psio[(i - 1), (jj - 1)] * bips.M((i - 1), (j - 1))
t[self._psi_o(j), 0] += val * (op.r + op.d)
for k in range(1, ngomod + 1):
for jj in range(1, ngomod + 1):
t[self._psi_o(j), self._psi_o(k)] -= M_psio[(i - 1), (jj - 1)] * bips.C((jj - 1), (j - 1),
(k - 1))
t = self.simplify_matrix(t)
tensor[self._psi_o(i)] = t
jacobian_tensor[self._psi_o(i)] = t + t.T
# deltaT_o part
for i in range(1, ngomod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
for jj in range(1, namod + 1):
t[0, 0] += bips.W((i - 1), (jj - 1)) * par.Cpgo[(jj - 1)]
for j in range(1, namod + 1):
val = 0
for jj in range(1, ngomod + 1):
val += U_inv[(i - 1), (jj - 1)] * bips.W((jj - 1), (j - 1))
t[self._theta_a(j), 0] += val * (2 * atp.sc * par.Lpgo + par.sbpa)
for j in range(1, ngomod + 1):
t[self._deltaT_o(j), 0] = - (par.Lpgo + par.sbpgo) * _kronecker_delta((i - 1), (j - 1))
for k in range(1, ngomod + 1):
for jj in range(1, ngomod + 1):
t[self._psi_o(j), self._deltaT_o(k)] -= U_inv[(i - 1), (jj - 1)] * bips.O((jj - 1), (j - 1),
(k - 1))
t = self.simplify_matrix(t)
tensor[self._deltaT_o(i)] = t
jacobian_tensor[self._deltaT_o(i)] = t + t.T
# deltaT_g part
if ground_temp:
for i in range(1, ngomod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
for jj in range(1, namod + 1):
t[0, 0] += bips.W[(i - 1), (jj - 1)] * par.Cpgo[(jj - 1)]
for j in range(1, namod + 1):
val = 0
for jj in range(1, ngomod + 1):
val += U_inv[(i - 1), (jj - 1)] * bips.W((jj - 1), (j - 1))
t[self._theta_a(j), 0] += val * (2 * atp.sc * par.Lpgo + par.sbpa)
for j in range(1, ngomod + 1):
t[self._deltaT_g(j), 0] = - (par.Lpgo + par.sbpgo) * _kronecker_delta((i - 1), (j - 1))
t = self.simplify_matrix(t)
tensor[self._deltaT_g(i)] = t
jacobian_tensor[self._deltaT_g(i)] = t + t.T
self.tensor = tensor.to_coo()
self.jacobian_tensor = jacobian_tensor.to_coo()
def compute_inner_products(self, num_threads=None):
"""Function computing and storing all the inner products at once.
Parameters
----------
num_threads: int or None, optional
Number of threads to use to compute the symbolic inner products. If `None` use all the cpus available.
Default to `None`.
"""
self._M = sp.zeros((self.noc, self.noc), dtype=float, format='dok')
self._U = sp.zeros((self.noc, self.noc), dtype=float, format='dok')
self._N = sp.zeros((self.noc, self.noc), dtype=float, format='dok')
self._O = sp.zeros((self.noc, self.noc, self.noc),
dtype=float,
format='dok')
self._C = sp.zeros((self.noc, self.noc, self.noc),
dtype=float,
format='dok')
if num_threads is None:
pool = multiprocessing.Pool(processes=multiprocessing.cpu_count())
else:
pool = multiprocessing.Pool(processes=num_threads)
# N inner products
l = [[(i, j), self.ip.ip_diff_x, (self._phi(i), self._phi(j))]
for i in range(self.noc) for j in range(self.noc)]
result = pool.map(_apply, l)
for res in result:
self._N[res[0]] = float(res[1].subs(self.subs).subs(
self.oceanic_basis.substitutions))
# M inner products
l = [[(i, j), self.ip.ip_lap, (self._phi(i), self._phi(j))]
for i in range(self.noc) for j in range(self.noc)]
result = pool.map(_apply, l)
for res in result:
self._M[res[0]] = float(res[1].subs(self.subs).subs(
self.oceanic_basis.substitutions))
# U inner products
l = [[(i, j), self.ip.symbolic_inner_product,
(self._phi(i), self._phi(j))] for i in range(self.noc)
for j in range(self.noc)]
result = pool.map(_apply, l)
for res in result:
self._U[res[0]] = float(res[1].subs(self.subs).subs(
self.oceanic_basis.substitutions))
# O inner products
l = [[(i, j, k), self.ip.ip_jac,
(self._phi(i), self._phi(j), self._phi(k))]
for i in range(self.noc) for j in range(self.noc)
for k in range(self.noc)]
result = pool.map(_apply, l)
for res in result:
self._O[res[0]] = float(res[1].subs(self.subs).subs(
self.oceanic_basis.substitutions))
# C inner products
l = [[(i, j, k), self.ip.ip_jac_lap,
(self._phi(i), self._phi(j), self._phi(k))]
for i in range(self.noc) for j in range(self.noc)
for k in range(self.noc)]
result = pool.map(_apply, l)
for res in result:
self._C[res[0]] = float(res[1].subs(self.subs).subs(
self.oceanic_basis.substitutions))
self._M = self._M.to_coo()
self._U = self._U.to_coo()
self._N = self._N.to_coo()
self._O = self._O.to_coo()
self._C = self._C.to_coo()
pool.terminate()
def compute_inner_products(self, num_threads=None):
"""Function computing and storing all the inner products at once.
Parameters
----------
num_threads: int or None, optional
Number of threads to use to compute the symbolic inner products. If `None` use all the cpus available.
Default to `None`.
"""
self._a = sp.zeros((self.natm, self.natm), dtype=float, format='dok')
self._u = sp.zeros((self.natm, self.natm), dtype=float, format='dok')
self._c = sp.zeros((self.natm, self.natm), dtype=float, format='dok')
self._b = sp.zeros((self.natm, self.natm, self.natm),
dtype=float,
format='dok')
self._g = sp.zeros((self.natm, self.natm, self.natm),
dtype=float,
format='dok')
if num_threads is None:
pool = multiprocessing.Pool(processes=multiprocessing.cpu_count())
else:
pool = multiprocessing.Pool(processes=num_threads)
# a inner products
args_list = [[(i, j), self.ip.ip_lap, (self._F(i), self._F(j))]
for i in range(self.natm) for j in range(self.natm)]
result = pool.map(_apply, args_list)
for res in result:
self._a[res[0]] = float(res[1].subs(self.subs).subs(
self.atmospheric_basis.substitutions))
# u inner products
args_list = [[(i, j), self.ip.symbolic_inner_product,
(self._F(i), self._F(j))] for i in range(self.natm)
for j in range(self.natm)]
result = pool.map(_apply, args_list)
for res in result:
self._u[res[0]] = float(res[1].subs(self.subs).subs(
self.atmospheric_basis.substitutions))
# c inner products
args_list = [[(i, j), self.ip.ip_diff_x, (self._F(i), self._F(j))]
for i in range(self.natm) for j in range(self.natm)]
result = pool.map(_apply, args_list)
for res in result:
self._c[res[0]] = float(res[1].subs(self.subs).subs(
self.atmospheric_basis.substitutions))
# b inner products
args_list = [[(i, j, k), self.ip.ip_jac_lap,
(self._F(i), self._F(j), self._F(k))]
for i in range(self.natm) for j in range(self.natm)
for k in range(self.natm)]
result = pool.map(_apply, args_list)
for res in result:
self._b[res[0]] = float(res[1].subs(self.subs).subs(
self.atmospheric_basis.substitutions))
# g inner products
args_list = [[(i, j, k), self.ip.ip_jac,
(self._F(i), self._F(j), self._F(k))]
for i in range(self.natm) for j in range(self.natm)
for k in range(self.natm)]
result = pool.map(_apply, args_list)
for res in result:
self._g[res[0]] = float(res[1].subs(self.subs).subs(
self.atmospheric_basis.substitutions))
self._a = self._a.to_coo()
self._u = self._u.to_coo()
self._c = self._c.to_coo()
self._g = self._g.to_coo()
self._b = self._b.to_coo()
pool.terminate()
def connect_to_ground(self,
ground_basis,
orographic_basis,
num_threads=None):
"""Connect the atmosphere to the ground.
Parameters
----------
ground_basis: SymbolicBasis or GroundSymbolicInnerProducts
Basis of function of the ground or a symbolic ground inner products object containing the basis.
orographic_basis: str
String to select which component basis modes to use to develop the orography in series.
Can be either 'atmospheric' or 'ground'.
num_threads: int or None, optional
Number of threads to use to compute the symbolic inner products. If `None` use all the cpus available.
Default to `None`.
"""
if isinstance(ground_basis, GroundSymbolicInnerProducts):
ground_basis = ground_basis.ground_basis
self.oceanic_basis = None
self.connected_to_ocean = False
self.ground_basis = ground_basis
self.connected_to_ground = True
if self.stored:
if num_threads is None:
pool = multiprocessing.Pool(
processes=multiprocessing.cpu_count())
else:
pool = multiprocessing.Pool(processes=num_threads)
ngr = len(ground_basis)
if orographic_basis == "atmospheric":
self._gh = None
else:
self._gh = sp.zeros((self.natm, self.natm, ngr),
dtype=float,
format='dok')
self._d = None
self._s = sp.zeros((self.natm, ngr), dtype=float, format='dok')
# s inner products
args_list = [[(i, j), self.iip.symbolic_inner_product,
(self._F(i), self._phi(j))] for i in range(self.natm)
for j in range(ngr)]
result = pool.map(_apply, args_list)
for res in result:
self._s[res[0]] = float(res[1].subs(self.subs).subs(
self.atmospheric_basis.substitutions).subs(
self.ground_basis.substitutions))
# gh inner products
args_list = [[(i, j, k), self.iip.ip_jac,
(self._F(i), self._F(j), self._phi(k))]
for i in range(self.natm) for j in range(self.natm)
for k in range(ngr)]
result = pool.map(_apply, args_list)
for res in result:
self._gh[res[0]] = float(res[1].subs(self.subs).subs(
self.atmospheric_basis.substitutions).subs(
self.ground_basis.substitutions))
self._s = self._s.to_coo()
if self._gh is not None:
self._gh = self._gh.to_coo()
pool.terminate()
def compute_tensor(self):
"""Routine to compute the tensor."""
aips = self.atmospheric_inner_products
oips = self.oceanic_inner_products
par = self.params
atp = par.atemperature_params
ap = par.atmospheric_params
op = par.oceanic_params
scp = par.scale_params
gp = par.ground_params
namod = par.nmod[0]
ngomod = par.nmod[1]
ndim = par.ndim
if par.gotemperature_params is not None:
ocean = par.gotemperature_params._name == "Oceanic Temperature"
ground_temp = par.gotemperature_params._name == "Ground Temperature"
else:
ocean = False
ground_temp = False
# 0-th tensor component is an empty matrix
tensor = sp.zeros((ndim + 1, ndim + 1, ndim + 1),
dtype=np.float64,
format='dok')
jacobian_tensor = sp.zeros((ndim + 1, ndim + 1, ndim + 1),
dtype=np.float64,
format='dok')
## Problem with matmul with object and DOK archi : Temporary fix until a better solution is found
hk = np.array(gp.hk, dtype=np.float)
g = aips.g.to_coo()
#################
# psi_a part
for i in range(1, namod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
for j in range(1, namod + 1):
t[self._psi_a(j), 0] = -((aips.c[(i - 1), (j - 1)] * scp.beta) / aips.a[(i - 1), (i - 1)]) \
- (ap.kd * _kronecker_delta((i - 1), (j - 1))) / 2
t[self._theta_a(j), 0] = (ap.kd * _kronecker_delta(
(i - 1), (j - 1))) / 2
if gp.hk is not None:
oro = (g[(i - 1), (j - 1), :] @ hk) / (2 * aips.a[(i - 1),
(i - 1)])
t[self._psi_a(j), 0] -= oro
t[self._theta_a(j), 0] += oro
for k in range(1, namod + 1):
t[self._psi_a(j), self._psi_a(k)] = - aips.b[(i - 1), (j - 1), (k - 1)] \
/ aips.a[(i - 1), (i - 1)]
t[self._theta_a(j), self._theta_a(k)] = - aips.b[(i - 1), (j - 1), (k - 1)] \
/ aips.a[(i - 1), (i - 1)]
if ocean:
for j in range(1, ngomod + 1):
t[self._psi_o(j), 0] = ap.kd * aips.d[(i - 1), (j - 1)] / \
(2 * aips.a[(i - 1), (i - 1)])
t = self.simplify_matrix(t)
tensor[self._psi_a(i)] = t
jacobian_tensor[self._psi_a(i)] = t + t.T
# theta_a part
for i in range(1, namod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
if par.Cpa is not None:
t[0, 0] = par.Cpa[i - 1] / (1 - aips.a[0, 0] * ap.sig0)
if atp.hd is not None and atp.thetas is not None:
t[0, 0] += atp.hd * atp.thetas[
(i - 1)] / (1. - ap.sig0 * aips.a[(i - 1), (i - 1)])
for j in range(1, namod + 1):
t[self._psi_a(j), 0] = (aips.a[(i - 1), (j - 1)] * ap.kd * ap.sig0) \
/ (-2 + 2 * aips.a[(i - 1), (i - 1)] * ap.sig0)
if par.LSBpa is not None and par.Lpa is not None:
heat = 2. * (par.LSBpa +
atp.sc * par.Lpa) * _kronecker_delta((i - 1),
(j - 1))
else:
heat = 0
t[self._theta_a(j),
0] = (-((ap.sig0 *
(2. * aips.c[(i - 1),
(j - 1)] * scp.beta + aips.a[(i - 1),
(j - 1)] *
(ap.kd + 4. * ap.kdp)))) +
heat) / (-2. + 2. * aips.a[(i - 1), (i - 1)] * ap.sig0)
if atp.hd is not None:
t[self._theta_a(j), 0] += (atp.hd * _kronecker_delta(
(i - 1), (j - 1))) / (ap.sig0 * aips.a[(i - 1),
(i - 1)] - 1.)
if gp.hk is not None:
oro = (ap.sig0 * g[(i - 1),
(j - 1), :] @ hk) / (2 * aips.a[
(i - 1), (i - 1)] * ap.sig0 - 2.)
t[self._theta_a(j), 0] -= oro
t[self._psi_a(j), 0] += oro
for k in range(1, namod + 1):
t[self._psi_a(j), self._theta_a(k)] = (aips.g[(i - 1), (j - 1), (k - 1)]
- aips.b[(i - 1), (j - 1), (k - 1)] * ap.sig0) / \
(-1 + aips.a[(i - 1), (i - 1)] * ap.sig0)
t[self._theta_a(j), self._psi_a(k)] = (aips.b[(i - 1), (j - 1), (k - 1)] * ap.sig0) \
/ (1 - aips.a[(i - 1), (i - 1)] * ap.sig0)
if ocean:
for j in range(1, ngomod + 1):
t[self._psi_o(j), 0] = ap.kd * (aips.d[(i - 1), (j - 1)] * ap.sig0) \
/ (2 - 2 * aips.a[(i - 1), (i - 1)] * ap.sig0)
if par.LSBpgo is not None and par.Lpa is not None:
t[self._deltaT_o(j), 0] = aips.s[(i - 1), (j - 1)] * (2 * par.LSBpgo + par.Lpa) \
/ (2 - 2 * aips.a[(i - 1), (i - 1)] * ap.sig0)
if ground_temp and i <= ngomod:
t[self._deltaT_g(i),
0] = (2 * par.LSBpgo +
par.Lpa) / (2 - 2 * aips.a[(i - 1), (i - 1)] * ap.sig0)
t = self.simplify_matrix(t)
tensor[self._theta_a(i)] = t
jacobian_tensor[self._theta_a(i)] = t + t.T
if ocean:
# psi_o part
for i in range(1, ngomod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
for j in range(1, namod + 1):
t[self._psi_a(j), 0] = oips.K[(i - 1), (j - 1)] * op.d \
/ (oips.M[(i - 1), (i - 1)] + par.G)
t[self._theta_a(j), 0] = -(oips.K[(i - 1), (j - 1)]) * op.d \
/ (oips.M[(i - 1), (i - 1)] + par.G)
for j in range(1, ngomod + 1):
t[self._psi_o(j),
0] = -((oips.N[(i - 1),
(j - 1)] * scp.beta + oips.M[(i - 1),
(i - 1)] *
(op.r + op.d) * _kronecker_delta(
(i - 1),
(j - 1)))) / (oips.M[(i - 1),
(i - 1)] + par.G)
for k in range(1, ngomod + 1):
t[self._psi_o(j), self._psi_o(k)] = -(oips.C[(i - 1), (j - 1), (k - 1)]) \
/ (oips.M[(i - 1), (i - 1)] + par.G)
t = self.simplify_matrix(t)
tensor[self._psi_o(i)] = t
jacobian_tensor[self._psi_o(i)] = t + t.T
# deltaT_o part
## Problem with matmul with object and DOK archi : Temporary fix until a better solution is found
Cpgo = np.array(par.Cpgo, dtype=np.float)
W = oips.W.to_coo()
#################
for i in range(1, ngomod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
t[0, 0] = W[(i - 1), :] @ Cpgo
for j in range(1, namod + 1):
t[self._theta_a(j),
0] = oips.W[(i - 1),
(j - 1)] * (2 * atp.sc * par.Lpgo + par.sbpa)
for j in range(1, ngomod + 1):
t[self._deltaT_o(j),
0] = -(par.Lpgo + par.sbpgo) * _kronecker_delta((i - 1),
(j - 1))
for k in range(1, ngomod + 1):
t[self._psi_o(j),
self._deltaT_o(k)] = -(oips.O[(i - 1), (j - 1),
(k - 1)])
t = self.simplify_matrix(t)
tensor[self._deltaT_o(i)] = t
jacobian_tensor[self._deltaT_o(i)] = t + t.T
# deltaT_g part
if ground_temp:
for i in range(1, ngomod + 1):
t = np.zeros((ndim + 1, ndim + 1), dtype=np.float64)
t[0, 0] = par.Cpgo[(i - 1)]
t[self._theta_a(i), 0] = 2 * atp.sc * par.Lpgo + par.sbpa
t[self._deltaT_g(i), 0] = -(par.Lpgo + par.sbpgo)
t = self.simplify_matrix(t)
tensor[self._deltaT_g(i)] = t
jacobian_tensor[self._deltaT_g(i)] = t + t.T
self.tensor = tensor.to_coo()
self.jacobian_tensor = jacobian_tensor.to_coo()
def empty_mask(mask_shape, dtype):
return sparse.zeros(mask_shape, dtype=dtype)
