feat: adds pdme fast calc for e field xs
This commit is contained in:
parent
45031857f2
commit
9e6d1df559
GPG Key ID:
BEBAEBF28083E022
@ -12,28 +12,44 @@ def fast_vs_for_dipoles(
|
||||
No error correction here baby.
|
||||
Expects dot_inputs to be numpy array of [rx, ry, rz, f] entries, so a n by 4 where n is number of measurement points.
|
||||
"""
|
||||
# name indexes:
|
||||
# A: dipole config index
|
||||
# cart: cartesian index
|
||||
# m: measurement index
|
||||
|
||||
# [A, cart]
|
||||
ps = dipoles[:, 0:3]
|
||||
ss = dipoles[:, 3:6]
|
||||
# [A]
|
||||
ws = dipoles[:, 6]
|
||||
|
||||
_logger.debug(f"ps: {ps}")
|
||||
_logger.debug(f"ss: {ss}")
|
||||
_logger.debug(f"ws: {ws}")
|
||||
|
||||
# [m, cart]
|
||||
rs = dot_inputs[:, 0:3]
|
||||
# [m]
|
||||
fs = dot_inputs[:, 3]
|
||||
|
||||
# [m, cart] - [A, 1, cart]
|
||||
# [A, m, cart]
|
||||
diffses = rs - ss[:, None]
|
||||
|
||||
_logger.debug(f"diffses: {diffses}")
|
||||
# [A, m]
|
||||
norms = numpy.linalg.norm(diffses, axis=2) ** 3
|
||||
_logger.debug(f"norms: {norms}")
|
||||
# [A, m, cart] [A, cart] -> [A, m]
|
||||
ases = (numpy.einsum("...ji, ...i", diffses, ps) / norms) ** 2
|
||||
_logger.debug(f"ases: {ases}")
|
||||
|
||||
# [A, 1], denom [m] + [A, 1] -> [A, m]
|
||||
# [A, m]
|
||||
bses = 2 * ws[:, None] / ((numpy.pi * fs) ** 2 + ws[:, None] ** 2)
|
||||
|
||||
_logger.debug(f"bses: {bses}")
|
||||
# returns shape [A, m]
|
||||
return ases * bses
|
||||
|
||||
|
||||
@ -70,6 +86,121 @@ def fast_vs_for_dipoleses(
|
||||
return numpy.einsum("...j->...", ases * bses)
|
||||
|
||||
|
||||
def fast_efieldxs_for_dipoles(
|
||||
dot_inputs: numpy.ndarray, dipoles: numpy.ndarray
|
||||
) -> numpy.ndarray:
|
||||
"""
|
||||
No error correction here baby.
|
||||
Expects dot_inputs to be numpy array of [rx, ry, rz, f] entries, so a n by 4 where n is number of measurement points.
|
||||
"""
|
||||
# name indexes:
|
||||
# A: dipole config index
|
||||
# j: dipole index within a config
|
||||
# cart: cartesian index
|
||||
# m: measurement index
|
||||
|
||||
# Indexes [A, cart]
|
||||
ps = dipoles[:, 0:3]
|
||||
ss = dipoles[:, 3:6]
|
||||
# Indexes [A]
|
||||
ws = dipoles[:, 6]
|
||||
|
||||
# Indexes [m, cart]
|
||||
rs = dot_inputs[:, 0:3]
|
||||
# Indexes [m]
|
||||
fs = dot_inputs[:, 3]
|
||||
|
||||
# Indexes [m, cart] - [A, 1, cart]
|
||||
# Broadcasting from right
|
||||
# diffses.indexes [A, m, cart]
|
||||
diffses = rs - ss[:, None, :]
|
||||
|
||||
# [A, m, cart][2] = cart
|
||||
# norms takes out axis 2, the last one, giving [A, m]
|
||||
norms = numpy.linalg.norm(diffses, axis=2)
|
||||
|
||||
# long story but this ends up becoming (A, 1, j)
|
||||
# Some evidence is looking at ps term, which has shape (A, 1, j, cart=0) becoming (A, 1, j)
|
||||
|
||||
# [A, m, cart] einsum [A, cart] explicitly gives [A, m]
|
||||
dot_products = numpy.einsum("amc,ac->am", diffses, ps) / (norms**2)
|
||||
|
||||
# [m, A] * [cart, m, A] -> [cart, m, A], transpose that and you get [A, m, cart]
|
||||
projections = numpy.transpose(
|
||||
numpy.transpose(dot_products) * numpy.transpose(diffses)
|
||||
)
|
||||
|
||||
# numerator [A, m, cart] - [A, 1, cart] -> [A, m, cart][:, :, 0] -> [A, m]
|
||||
alphas = (3 * projections - ps[:, numpy.newaxis])[:, :, 0] / norms**3
|
||||
|
||||
# [A, m]
|
||||
ases = alphas**2
|
||||
|
||||
# [A, 1], denom [m] + [A, 1] -> [A, m]
|
||||
# [A, m]
|
||||
bses = 2 * ws[:, None] / ((numpy.pi * fs) ** 2 + ws[:, None] ** 2)
|
||||
|
||||
# return shape [A, m, j]
|
||||
return ases * bses
|
||||
|
||||
|
||||
def fast_efieldxs_for_dipoleses(
|
||||
dot_inputs: numpy.ndarray, dipoleses: numpy.ndarray
|
||||
) -> numpy.ndarray:
|
||||
"""
|
||||
No error correction here baby.
|
||||
Expects dot_inputs to be numpy array of [rx, ry, rz, f] entries, so a n by 4 where n is number of measurement points.
|
||||
|
||||
Dipoleses are expected to be array of arrays of arrays:
|
||||
list of sets of dipoles which are part of a single arrangement to be added together.
|
||||
"""
|
||||
# name indexes:
|
||||
# A: dipole config index
|
||||
# j: dipole index within a config
|
||||
# cart: cartesian index
|
||||
# m: measurement index
|
||||
|
||||
# Indexes [A, j, cart]
|
||||
ps = dipoleses[:, :, 0:3]
|
||||
ss = dipoleses[:, :, 3:6]
|
||||
# Indexes [A, j]
|
||||
ws = dipoleses[:, :, 6]
|
||||
|
||||
# Indexes [m, cart]
|
||||
rs = dot_inputs[:, 0:3]
|
||||
# Indexes [m]
|
||||
fs = dot_inputs[:, 3]
|
||||
|
||||
# Indexes [m, 1, cart] - [A, 1, j, cart]
|
||||
# Broadcasting from right
|
||||
# diffses.indexes [A, m, j, cart]
|
||||
diffses = rs[:, None] - ss[:, None, :]
|
||||
|
||||
# norms takes out axis 3, the last one, giving [A, m, j]
|
||||
norms = numpy.linalg.norm(diffses, axis=3)
|
||||
|
||||
# long story but this ends up becoming (A, 1, j)
|
||||
# Some evidence is looking at ps term, which has shape (A, 1, j, cart=0) becoming (A, 1, j)
|
||||
alphas = (
|
||||
(
|
||||
3
|
||||
* numpy.transpose(
|
||||
numpy.transpose(
|
||||
numpy.einsum("abcd,acd->abc", diffses, ps) / (norms**2)
|
||||
)
|
||||
* numpy.transpose(diffses)
|
||||
)[:, :, :, 0]
|
||||
)
|
||||
- ps[:, numpy.newaxis, :, 0]
|
||||
) / (norms**3)
|
||||
ases = alphas**2
|
||||
# bses.shape [A, m, j)]
|
||||
bses = 2 * ws[:, None, :] / ((numpy.pi * fs[:, None]) ** 2 + ws[:, None, :] ** 2)
|
||||
|
||||
# return shape [A, m, j]
|
||||
return numpy.einsum("...j->...", ases * bses)
|
||||
|
||||
|
||||
def fast_vs_for_asymmetric_dipoleses(
|
||||
dot_inputs: numpy.ndarray, dipoleses: numpy.ndarray, temp: numpy.ndarray
|
||||
) -> numpy.ndarray:
|
||||
|
||||
@ -17,6 +17,15 @@ def s_potential_from_arrays(
|
||||
return dipole.s_electric_potential_at_position(r, f)
|
||||
|
||||
|
||||
def s_electric_field_x_from_arrays(
|
||||
dipole_array: numpy.ndarray, dotf_array: numpy.ndarray
|
||||
) -> float:
|
||||
dipole = dipole_from_array(dipole_array)
|
||||
r = dotf_array[0:3]
|
||||
f = dotf_array[3]
|
||||
return dipole.s_electric_fieldx_at_position(r, f)
|
||||
|
||||
|
||||
def test_fast_v_calc():
|
||||
d1 = [1, 2, 3, 4, 5, 6, 7]
|
||||
d2 = [2, 5, 3, 4, -5, -6, 2]
|
||||
@ -107,6 +116,96 @@ def test_fast_v_calc_big_multidipole():
|
||||
)
|
||||
|
||||
|
||||
def test_fast_electric_field_x_calc():
|
||||
d1 = [1, 2, 3, 4, 5, 6, 7]
|
||||
d2 = [2, 5, 3, 4, -5, -6, 2]
|
||||
|
||||
dipoles = numpy.array([d1, d2])
|
||||
|
||||
dot_inputs = numpy.array([[-1, -1, -1, 11], [2, 3, 1, 5.5]])
|
||||
|
||||
expected_11 = s_electric_field_x_from_arrays(dipoles[0], dot_inputs[0])
|
||||
expected_12 = s_electric_field_x_from_arrays(dipoles[1], dot_inputs[0])
|
||||
expected_21 = s_electric_field_x_from_arrays(dipoles[0], dot_inputs[1])
|
||||
expected_22 = s_electric_field_x_from_arrays(dipoles[1], dot_inputs[1])
|
||||
|
||||
expected = numpy.array([[expected_11, expected_21], [expected_12, expected_22]])
|
||||
|
||||
numpy.testing.assert_allclose(
|
||||
pdme.util.fast_v_calc.fast_efieldxs_for_dipoles(dot_inputs, dipoles),
|
||||
expected,
|
||||
err_msg="E x fast calc at dot aren't as expected.",
|
||||
)
|
||||
|
||||
|
||||
def test_fast_electric_field_x_calc_multidipoles():
|
||||
d1 = [1, 2, 3, 4, 5, 6, 7]
|
||||
d2 = [2, 5, 3, 4, -5, -6, 2]
|
||||
|
||||
dipoles = numpy.array([[d1, d2]])
|
||||
|
||||
dot_inputs = numpy.array([[-1, -1, -1, 11], [2, 3, 1, 5.5]])
|
||||
|
||||
expected_11 = s_electric_field_x_from_arrays(dipoles[0][0], dot_inputs[0])
|
||||
expected_12 = s_electric_field_x_from_arrays(dipoles[0][1], dot_inputs[0])
|
||||
expected_21 = s_electric_field_x_from_arrays(dipoles[0][0], dot_inputs[1])
|
||||
expected_22 = s_electric_field_x_from_arrays(dipoles[0][1], dot_inputs[1])
|
||||
|
||||
expected = numpy.array([[expected_11 + expected_12, expected_21 + expected_22]])
|
||||
|
||||
numpy.testing.assert_allclose(
|
||||
pdme.util.fast_v_calc.fast_efieldxs_for_dipoleses(dot_inputs, dipoles),
|
||||
expected,
|
||||
err_msg="E x fast calc at dot aren't as expected for multidipole calc.",
|
||||
)
|
||||
|
||||
|
||||
def test_fast_electric_field_x_calc_big_multidipole():
|
||||
|
||||
dipoleses = numpy.array(
|
||||
[
|
||||
[
|
||||
[1, 1, 5, 6, 3, 1, 1],
|
||||
[5, 3, 2, 13, 1, 1, 2],
|
||||
[-5, -5, -3, -1, -3, 8, 3],
|
||||
],
|
||||
[
|
||||
[-3, -1, -2, -2, -6, 3, 4],
|
||||
[8, 0, 2, 0, 1, 5, 5],
|
||||
[1, 4, -4, -1, -3, -5, 6],
|
||||
],
|
||||
]
|
||||
)
|
||||
|
||||
dot_inputs = numpy.array(
|
||||
[
|
||||
[1, 1, 0, 1],
|
||||
[2, 5, 6, 2],
|
||||
[3, 1, 3, 3],
|
||||
[0.5, 0.5, 0.5, 4],
|
||||
]
|
||||
)
|
||||
|
||||
expected = [
|
||||
[
|
||||
sum(
|
||||
[
|
||||
s_electric_field_x_from_arrays(dipole_array, dot_input)
|
||||
for dipole_array in dipole_config
|
||||
]
|
||||
)
|
||||
for dot_input in dot_inputs
|
||||
]
|
||||
for dipole_config in dipoleses
|
||||
]
|
||||
|
||||
numpy.testing.assert_allclose(
|
||||
pdme.util.fast_v_calc.fast_efieldxs_for_dipoleses(dot_inputs, dipoleses),
|
||||
expected,
|
||||
err_msg="E x fast calc at dot aren't as expected for multidipole calc.",
|
||||
)
|
||||
|
||||
|
||||
def test_between():
|
||||
low = numpy.array([1, 2, 3])
|
||||
high = numpy.array([6, 7, 8])
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user