Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Reduce memory usage for timeseries jac computation #1001

Merged
merged 5 commits into from
Oct 16, 2023
Merged

Reduce memory usage for timeseries jac computation #1001

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
06607da
jac fix for timeseries derivs
johnjasa Oct 12, 2023
074bde7
Also updated analytic phase deriv calc
johnjasa Oct 12, 2023
04530a9
Removed use_sparse flag
johnjasa Oct 16, 2023
8bdf903
Merging
johnjasa Oct 16, 2023
fba061d
pycodestyle updates
johnjasa Oct 16, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 27 additions & 19 deletions dymos/transcriptions/analytic/analytic_timeseries_output_comp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,6 @@
from ..common.timeseries_output_comp import TimeseriesOutputCompBase


_USE_SPARSE = True


class AnalyticTimeseriesOutputComp(TimeseriesOutputCompBase):
"""
Class definition of the timeseries output comp for AnalyticPhase.
Expand Down Expand Up @@ -86,12 +83,8 @@ def setup(self):
L_blocks.append(L)
D_blocks.append(D)

if _USE_SPARSE:
self.interpolation_matrix = sp.block_diag(L_blocks, format='csr')
self.differentiation_matrix = sp.block_diag(D_blocks, format='csr')
else:
self.interpolation_matrix = block_diag(*L_blocks)
self.differentiation_matrix = block_diag(*D_blocks)
self.interpolation_matrix = sp.block_diag(L_blocks, format='csr')
self.differentiation_matrix = sp.block_diag(D_blocks, format='csr')

self.add_input('dt_dstau', shape=(self.input_num_nodes,), units=self.options['time_units'])

Expand Down Expand Up @@ -150,35 +143,50 @@ def _add_output_configure(self, name, units, shape, desc='', src=None, rate=Fals
self._vars[name] = (input_name, output_name, shape, rate)

size = np.prod(shape)
jac = np.zeros((output_num_nodes, size, input_num_nodes, size))

if rate:
mat = self.differentiation_matrix
else:
mat = self.interpolation_matrix

for i in range(size):
if _USE_SPARSE:
jac[:, i, :, i] = mat.toarray()
else:
jac[:, i, :, i] = mat
# Preallocate lists to hold data for jac
jac_data = []
jac_indices = []
jac_indptr = [0]

# Iterate over the dense dimension 'size'
for s in range(size):
# Extend the data and indices using the CSR attributes of mat
jac_data.extend(mat.data)
jac_indices.extend(mat.indices + s * input_num_nodes)

# For every non-zero row in mat, update jac's indptr
new_indptr = mat.indptr[1:] + s * len(mat.data)
jac_indptr.extend(new_indptr)

# Correct the last entry of jac_indptr
jac_indptr[-1] = len(jac_data)

jac = jac.reshape((output_num_nodes * size, input_num_nodes * size), order='C')
# Construct the sparse jac matrix in CSR format
jac = sp.csr_matrix((jac_data, jac_indices, jac_indptr),
shape=(output_num_nodes * size, input_num_nodes * size))

jac_rows, jac_cols = np.nonzero(jac)
# Now, if you want to get the row and column indices of the non-zero entries in the jac matrix:
jac_rows, jac_cols = jac.nonzero()

# There's a chance that the input for this output was pulled from another variable with
# different units, so account for that with a conversion.
val = np.squeeze(np.array(jac[jac_rows, jac_cols]))
if input_units is None or units is None:
self.declare_partials(of=output_name, wrt=input_name,
rows=jac_rows, cols=jac_cols, val=jac[jac_rows, jac_cols])
rows=jac_rows, cols=jac_cols, val=val)
else:
scale, offset = unit_conversion(input_units, units)
self._conversion_factors[output_name] = scale, offset

self.declare_partials(of=output_name, wrt=input_name,
rows=jac_rows, cols=jac_cols,
val=scale * jac[jac_rows, jac_cols])
val=scale * val)

return added_source

Expand Down
46 changes: 27 additions & 19 deletions dymos/transcriptions/pseudospectral/components/pseudospectral_timeseries_output_comp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,6 @@
from ...common.timeseries_output_comp import TimeseriesOutputCompBase


_USE_SPARSE = True


class PseudospectralTimeseriesOutputComp(TimeseriesOutputCompBase):
"""
Class definition of the PseudospectralTimeseriesOutputComp.
Expand Down Expand Up @@ -86,12 +83,8 @@ def setup(self):
L_blocks.append(L)
D_blocks.append(D)

if _USE_SPARSE:
self.interpolation_matrix = sp.block_diag(L_blocks, format='csr')
self.differentiation_matrix = sp.block_diag(D_blocks, format='csr')
else:
self.interpolation_matrix = block_diag(*L_blocks)
self.differentiation_matrix = block_diag(*D_blocks)
self.interpolation_matrix = sp.block_diag(L_blocks, format='csr')
self.differentiation_matrix = sp.block_diag(D_blocks, format='csr')

self.add_input('dt_dstau', shape=(self.input_num_nodes,), units=self.options['time_units'])

Expand Down Expand Up @@ -150,35 +143,50 @@ def _add_output_configure(self, name, units, shape, desc='', src=None, rate=Fals
self._vars[name] = (input_name, output_name, shape, rate)

size = np.prod(shape)
jac = np.zeros((output_num_nodes, size, input_num_nodes, size))

if rate:
mat = self.differentiation_matrix
else:
mat = self.interpolation_matrix

for i in range(size):
if _USE_SPARSE:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The _USE_SPARSE setting was just used to quickly toggle and check performance between sparse and dense implementations. Let's remove the _USE_SPARSE variable and just assume that we always do so. Any performance benefits of dense are typically only present in smaller problems.

jac[:, i, :, i] = mat.toarray()
else:
jac[:, i, :, i] = mat
# Preallocate lists to hold data for jac
jac_data = []
jac_indices = []
jac_indptr = [0]

# Iterate over the dense dimension 'size'
for s in range(size):
# Extend the data and indices using the CSR attributes of mat
jac_data.extend(mat.data)
jac_indices.extend(mat.indices + s * input_num_nodes)

# For every non-zero row in mat, update jac's indptr
new_indptr = mat.indptr[1:] + s * len(mat.data)
jac_indptr.extend(new_indptr)

# Correct the last entry of jac_indptr
jac_indptr[-1] = len(jac_data)

jac = jac.reshape((output_num_nodes * size, input_num_nodes * size), order='C')
# Construct the sparse jac matrix in CSR format
jac = sp.csr_matrix((jac_data, jac_indices, jac_indptr),
shape=(output_num_nodes * size, input_num_nodes * size))

jac_rows, jac_cols = np.nonzero(jac)
# Now, if you want to get the row and column indices of the non-zero entries in the jac matrix:
jac_rows, jac_cols = jac.nonzero()

# There's a chance that the input for this output was pulled from another variable with
# different units, so account for that with a conversion.
val = np.squeeze(np.array(jac[jac_rows, jac_cols]))
if input_units is None or units is None:
self.declare_partials(of=output_name, wrt=input_name,
rows=jac_rows, cols=jac_cols, val=jac[jac_rows, jac_cols])
rows=jac_rows, cols=jac_cols, val=val)
else:
scale, offset = unit_conversion(input_units, units)
self._conversion_factors[output_name] = scale, offset

self.declare_partials(of=output_name, wrt=input_name,
rows=jac_rows, cols=jac_cols,
val=scale * jac[jac_rows, jac_cols])
val=scale * val)

return added_source

Expand Down