diff --git a/pdme/util/fast_v_calc.py b/pdme/util/fast_v_calc.py index 0f8351c..0fd7488 100644 --- a/pdme/util/fast_v_calc.py +++ b/pdme/util/fast_v_calc.py @@ -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: diff --git a/tests/util/test_fast_v_calc.py b/tests/util/test_fast_v_calc.py index b14b891..d475729 100644 --- a/tests/util/test_fast_v_calc.py +++ b/tests/util/test_fast_v_calc.py @@ -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])