def reduce_einstein_coefficients_slow(a_info_nnz, energy_levels):
"""
Construct the A_matrix from the nnz info of the Einstein coefficients
Use the nonzero entries of the Einstein coefficients to construct the
A matrix for the levels found in the energy_levels parameter.
:param DataSetRawBase.A_info_nnz like a_info_nnz: a tuple of elements v_nnz,
j_nnz, vp_nnz, jp_nnz, A_nnz
:param energy_levels: numpy array array of N rows, the v and j levels and
their corresponding energy. The elements of the array are stored in
increasing order in the energy. This array can be obtained e.g. from
read_levels.read_levels_lique()
:return: A 2D matrix of shape (energy_levels.size, energy_levels.size) that
is a lower triangular matrix containing the Einstein coefficients.
"""
check_self_transitions_in_einstien_nnz_data(a_info_nnz)
v_nnz, j_nnz, vp_nnz, jp_nnz, a_nnz = a_info_nnz
levels = energy_levels
# get the unique label for the (v,j) pairs
labels_ini = linear_2d_index(v_nnz, j_nnz, n_i=levels.v_max_allowed)
labels_fin = linear_2d_index(vp_nnz, jp_nnz, n_i=levels.v_max_allowed)
a_reduced = zeros((levels.size, levels.size), 'f8')
for i, A_i in enumerate(a_nnz):
# print('{:4}/{:4}'.format(i+1, len(A_nnz)))
# get the indices based on v,j, jp, jp comaprisons
# v, j, vp, jp = v_nnz[i], j_nnz[i], vp_nnz[i], jp_nnz[i]
# ind_ini = where((levels['v'] == v)*(levels['j'] == j))[0]
# ind_fin = where((levels['v'] == vp)*(levels['j'] == jp))[0]
# get the indices based on label comparisons
ind_ini = where(levels.data['label'] == labels_ini[i])[0]
ind_fin = where(levels.data['label'] == labels_fin[i])[0]
if ind_ini.size != 0 or ind_fin.size != 0:
a_reduced[ind_ini, ind_fin] = A_i
else:
continue
# DEBUG
# A_reduced[A_reduced > 0.0] = numpy.log10(A_reduced[A_reduced > 0.0])
# pylab.imshow(A_reduced, interpolation='none')
# pylab.colorbar()
# pylab.show()
return a_reduced
def reduce_einstein_coefficients(a_mat, energy_levels):
"""
Recude a 4D A transition ndarray into a 2D matrix
.. todo:: rename this function to
reduce_einstein_coeffients_two_quantum_numbers or a name that indicates a
molecule whose levels are identified by two "quantum" numbers or labels.
Given the array A that is indexed using four indices A[v, j, v', j']
returns an array A_reduced that is indexed with two indices A_reduced[i, f]
where i and f are the initial and final levels respectively.
:param ndarray a_mat: A 4D matrix holding the A coefficients.
A[v, j, v', j'] where (v,j) are the initial levels and (v',j') are the
final levels.
:param EnergyLevelsSpeciesBase energy_levels: numpy array array of N rows,
the v and j levels and their corresponding energy. The elements of the
array are stored in increasing order in the energy.
:return: A_reduced that is a square matrix of shape
(energy_levels.size, energy_levels.size) with the nonzero values of A
mapped to the indices of the levels in the array energy_levels.
"""
levels = energy_levels
n_levels = len(energy_levels.data)
labels = energy_levels.data['label']
# find the non zeros elements and their corresponding indices in A
(v_nnz, j_nnz, vp_nnz, jp_nnz), a_nnz = where(a_mat > 0), a_mat[a_mat > 0]
# get the unique label for the (v,j) pairs
labels_ini = linear_2d_index(v_nnz, j_nnz, n_i=levels.v_max_allowed)
labels_fin = linear_2d_index(vp_nnz, jp_nnz, n_i=levels.v_max_allowed)
# keep transitions whose initial levels labels and the final label of the
# transition are found in energy_levels
mask = in1d(labels_ini, labels) * in1d(labels_fin, labels)
labels_ini, labels_fin = labels_ini[mask], labels_fin[mask]
a_nnz = a_nnz[mask]
# get the indices of the labels of the levels in the transitions (that are
# now all a subset of the energy_levels)
inds_ini = find_matching_indices(labels, labels_ini)
inds_fin = find_matching_indices(labels, labels_fin)
# define the reduced A matrix and fill it up using inds_ini and inds_fin
a_reduced = zeros((n_levels, n_levels), 'f8') * a_nnz.unit
a_reduced[inds_ini, inds_fin] = a_nnz
# DEBUG
# display_matrix(a_reduced.value)
return a_reduced
def test_that_linear_2d_index_has_no_clashes():
n = 1000000
i, j = numpy.unique(
numpy.vstack(
(numpy.random.randint(0, 100, n),
numpy.random.randint(0, 100, n))
), axis=1
)
linear_inds = linear_2d_index(i, j)
assert linear_inds.size == numpy.unique(linear_inds).size
def set_labels(self, v_max=None):
"""
set the field self.data['label']. This method modifies
self.data['label'] and self.v_max_allowed
:param v_max: The maximum value of v to be used in computing and
setting the labels.
"""
if v_max is not None:
self.v_max_allowed = v_max
self.data['label'] = utils.linear_2d_index(self.data['v'],
self.data['j'],
n_i=v_max)
def reduce_collisional_coefficients_slow(
cr_info_nnz,
energy_levels,
set_inelastic_coefficient_to_zero=False,
set_excitation_coefficients_to_zero=False,
reduced_data_is_upper_to_lower_only=True):
"""
Construct the K_matrix from the sparese nnz collisional coefficeints data
Given the data that is available for the collisional transitions as a
function of initial and final v,j -> v',j' levels and for different values
of temperature, returns an array of the data reduced as a matrix for each
temperature value and mapped by the energy level labels instead of v and j.
:param tuple cr_info_nnz: The information of the collisional data of the
levels for which data is available. This parameter is usually returned
by a read of raw collisional data e.g
read_collision_coefficients_lique_and_wrathmall.
The tuple has four elements:
- (v,j): The first element is a 2D array of shape (2, n_transitions)
that are the v and j of the initial level of the transitions (ini).
The columns of this array (v, j = ini) are the initial v and j of
the transitions for a certain T section. i.e. the number of non-zero
elements in K for a certain temperature is equal to the number of
elements in the v or j columns.
v.size = j.size = where(K[..., 0] > 0)[0].size
- (v',j') The second element is the same of the first element but for
the final level of the transitions (v', j' = fin)
- unique_nnz: The third element is a 2D array of shape
(2, n_unique_transitions) that are the unique levels involved in all
the transitions.
- cr_nnz: The last element is an array of shape (T.size, n_transitions)
which are the collisional coefficient rates with non-zero values
for each value of temperature in the T array.
:param EnergyLevelsMolecular energy_levels: The energy levels object whose
evergy levels map to the collisional nnz data.
:param bool set_inelastic_coefficient_to_zero: If True, then all the
elements along the diagonal of the matrix for all the temperatures are set
to zero. i.e the collision coefficient for the in-elastic collisions
would be ignored.
:param bool set_excitation_coefficients_to_zero: If True, then all the
elements in the upper triangular part of the matrix for all the
temperatures are set to zero. i.e the lower to upper transitions
(excitation) are ignored.
:param bool reduced_data_is_upper_to_lower_only: If True, then it is assumed
that the computed reduced matrix has only upper to lower
(i.e de-excitation) coefficients. If there is any non-zero value along the
diagonal (i.e in-elastic coefficients) or any non zero value in the upper
triangular part (i.e excitation coeffients), then an exception is raised.
:return: The reduced matrices of the collisional coefficients. One matrix
for each temperature value. The shape of the matrix is
(n_levels, n_levels, n_temperature_values)
"""
# check_self_transitions_in_einstien_nnz_data(A_info_nnz)
levels = energy_levels
n_levels = len(energy_levels.data)
labels = energy_levels.data['label']
(v_nnz, j_nnz), (vp_nnz, jp_nnz), unique_nnz, cr_nnz = cr_info_nnz
# get the unique label for the (v,j) pairs
labels_ini = linear_2d_index(v_nnz, j_nnz, n_i=levels.v_max_allowed)
labels_fin = linear_2d_index(vp_nnz, jp_nnz, n_i=levels.v_max_allowed)
# number of temperature value for which collisional data is available
n_T = cr_nnz.shape[0]
k_dex_reduced = zeros((n_levels, n_levels, n_T), 'f8') * cr_nnz.unit
for i, cr_i in enumerate(cr_nnz.T):
# print('{:4}/{:4}'.format(i+1, len(A_nnz)))
# get the indices based on v,j, jp, jp comparisons
# v, j, vp, jp = v_nnz[i], j_nnz[i], vp_nnz[i], jp_nnz[i]
# ind_ini = where((levels['v'] == v)*(levels['j'] == j))[0]
# ind_fin = where((levels['v'] == vp)*(levels['j'] == jp))[0]
# get the indices based on label comparisons
ind_ini = where(labels == labels_ini[i])[0]
ind_fin = where(labels == labels_fin[i])[0]
if ind_ini.size != 0 or ind_fin.size != 0:
k_dex_reduced[ind_ini, ind_fin, :] = cr_i
else:
continue
#
# optionally zero out data above the diagonals
#
if set_inelastic_coefficient_to_zero:
i_diag, j_diag = numpy.diag_indices(n_levels)
k_dex_reduced[i_diag, j_diag, :] = 0.0
if set_excitation_coefficients_to_zero:
i_upper, j_upper = numpy.triu_indices(n_levels, 1)
k_dex_reduced[i_upper, j_upper, :] = 0.0
if reduced_data_is_upper_to_lower_only:
# check that the upper triangular matrices for all the temperatures
# including the diagonal are zero since this is the K_dex matrix
# (see doc)
assert numpy.triu(numpy.moveaxis(k_dex_reduced, -1, 0)).sum() == 0.0
# DEBUG
# K_dex_reduced[K_dex_reduced > 0.0] = numpy.log10(
# K_dex_reduced[K_dex_reduced > 0.0])
# pylab.imshow(K_dex_reduced[:, :, 0], interpolation='none')
# pylab.colorbar()
# pylab.show()
return k_dex_reduced