Merge pull request #222 from CovertLab/optimizations

ecoli/models/polypeptide_elongation_models.py

@@ -2,6 +2,7 @@
 
 from typing import Callable, Optional, Tuple
 
+from numba import njit
 import numpy as np
 from scipy.integrate import solve_ivp
 from vivarium.library.units import units as vivunits
@@ -565,7 +566,8 @@ def isTimeStepShortEnough(self, inputTimeStep, timeStepSafetyFraction):
             short_enough = False
 
         return short_enough
-
+
+
 def ppgpp_metabolite_changes(uncharged_trna_conc, charged_trna_conc,
         ribosome_conc, f, rela_conc, spot_conc, ppgpp_conc, counts_to_molar,
         v_rib, charging_params, ppgpp_params, time_step,
@@ -700,6 +702,7 @@ def ppgpp_metabolite_changes(uncharged_trna_conc, charged_trna_conc,
 
     return delta_metabolites, n_syn_reactions, n_deg_reactions, v_rela_syn, v_spot_syn, v_deg, v_deg_inhibited
 
+
 def calculate_trna_charging(synthetase_conc, uncharged_trna_conc, charged_trna_conc, aa_conc, ribosome_conc,
         f, params, supply=None, time_limit=1000, limit_v_rib=False, use_disabled_aas=False):
     '''
@@ -762,34 +765,17 @@ def dcdt(t, c):
             ndarray[float]: dc/dt for tRNA concentrations
                 dims: 2 * number of amino acids (uncharged tRNA come first, then charged)
         '''
+        v_charging, dtrna, daa = dcdt_jit(t, c, n_aas_masked, n_aas, mask,
+            params['kS'], synthetase_conc, params['KMaa'], params['KMtf'],
+            f, params['krta'], params['krtf'], params['max_elong_rate'],
+            ribosome_conc, limit_v_rib, aa_rate_limit, v_rib_max)
 
-        uncharged_trna_conc = c[:n_aas_masked]
-        charged_trna_conc = c[n_aas_masked:2*n_aas_masked]
-        aa_conc = c[2*n_aas_masked:2*n_aas_masked+n_aas]
-        masked_aa_conc = aa_conc[mask]
-
-        v_charging = (params['kS'] * synthetase_conc * uncharged_trna_conc * masked_aa_conc / (params['KMaa'][mask] * params['KMtf'][mask])
-            / (1 + uncharged_trna_conc/params['KMtf'][mask] + masked_aa_conc/params['KMaa'][mask] + uncharged_trna_conc*masked_aa_conc/params['KMtf'][mask]/params['KMaa'][mask]))
-        with np.errstate(divide='ignore'):
-            numerator_ribosome = 1 + np.sum(f * (params['krta'] / charged_trna_conc + uncharged_trna_conc / charged_trna_conc * params['krta'] / params['krtf']))
-        v_rib = params['max_elong_rate'] * ribosome_conc / numerator_ribosome
-
-        # Handle case when f is 0 and charged_trna_conc is 0
-        if not np.isfinite(v_rib):
-            v_rib = 0
-
-        # Limit v_rib and v_charging to the amount of available amino acids
-        if limit_v_rib:
-            v_charging = np.fmin(v_charging, aa_rate_limit)
-            v_rib = min(v_rib, v_rib_max)
-
-        dtrna = v_charging - v_rib*f
-        daa = np.zeros(n_aas)
         if supply is None:
             v_synthesis = np.zeros(n_aas)
             v_import = np.zeros(n_aas)
             v_export = np.zeros(n_aas)
         else:
+            aa_conc = c[2*n_aas_masked:2*n_aas_masked+n_aas]
             v_synthesis, v_import, v_export = supply(unit_conversion * aa_conc)
             v_supply = v_synthesis + v_import - v_export
             daa[mask] = v_supply[mask] - v_charging
@@ -854,6 +840,38 @@ def dcdt(t, c):
 
     return new_fraction_charged, v_rib, total_synthesis, total_import, total_export
 
+
+@njit(error_model='numpy')
+def dcdt_jit(t, c, n_aas_masked, n_aas, mask,
+    kS, synthetase_conc, KMaa, KMtf,
+    f, krta, krtf, max_elong_rate,
+    ribosome_conc, limit_v_rib, aa_rate_limit, v_rib_max
+):
+    uncharged_trna_conc = c[:n_aas_masked]
+    charged_trna_conc = c[n_aas_masked:2*n_aas_masked]
+    aa_conc = c[2*n_aas_masked:2*n_aas_masked+n_aas]
+    masked_aa_conc = aa_conc[mask]
+
+    v_charging = (kS * synthetase_conc * uncharged_trna_conc * masked_aa_conc / (KMaa[mask] * KMtf[mask])
+        / (1 + uncharged_trna_conc/KMtf[mask] + masked_aa_conc/KMaa[mask] + uncharged_trna_conc*masked_aa_conc/KMtf[mask]/KMaa[mask]))
+    numerator_ribosome = 1 + np.sum(f * (krta / charged_trna_conc + uncharged_trna_conc / charged_trna_conc * krta / krtf))
+    v_rib = max_elong_rate * ribosome_conc / numerator_ribosome
+
+    # Handle case when f is 0 and charged_trna_conc is 0
+    if not np.isfinite(v_rib):
+        v_rib = 0
+
+    # Limit v_rib and v_charging to the amount of available amino acids
+    if limit_v_rib:
+        v_charging = np.fmin(v_charging, aa_rate_limit)
+        v_rib = min(v_rib, v_rib_max)
+
+    dtrna = v_charging - v_rib*f
+    daa = np.zeros(n_aas)
+
+    return v_charging, dtrna, daa
+
+
 def get_charging_supply_function(
         supply_in_charging: bool,
         mechanistic_supply: bool,

reconstruction/ecoli/dataclasses/process/metabolism.py

@@ -13,6 +13,7 @@
 import re
 from typing import Any, cast, Dict, Iterable, List, Optional, Set, Tuple, Union
 
+from numba import njit
 import numpy as np
 import sympy as sp
 from sympy.parsing.sympy_parser import parse_expr
@@ -1318,22 +1319,10 @@ def amino_acid_synthesis(self, counts_per_aa_fwd: np.ndarray, counts_per_aa_rev:
 		if units.hasUnit(aa_conc):
 			aa_conc = aa_conc.asNumber(METABOLITE_CONCENTRATION_UNITS)
 
-		km_saturation = np.product(1 / (1 + self.aa_upstream_kms / aa_conc), axis=1)
-
-		# Determine saturation fraction for reactions
-		forward_fraction = 1 / (1 + aa_conc / self.aa_kis) * km_saturation
-		reverse_fraction = 1 / (1 + self.aa_reverse_kms / aa_conc)
-		deg_fraction = 1 / (1 + self.aa_degradation_kms / aa_conc)
-		loss_fraction = reverse_fraction + deg_fraction
-
-		# Calculate synthesis rate
-		synthesis = (
-			self.aa_forward_stoich @ (self.aa_kcats_fwd * counts_per_aa_fwd * forward_fraction)
-			- self.aa_reverse_stoich @ (self.aa_kcats_rev * counts_per_aa_rev * reverse_fraction)
-			- self.aa_kcats_rev * counts_per_aa_rev * deg_fraction
-		)
-
-		return synthesis, forward_fraction, loss_fraction
+		return amino_acid_synthesis_jit(counts_per_aa_fwd, counts_per_aa_rev, aa_conc,
+			self.aa_upstream_kms, self.aa_kis, self.aa_reverse_kms,
+			self.aa_degradation_kms, self.aa_forward_stoich, self.aa_kcats_fwd,
+			self.aa_reverse_stoich, self.aa_kcats_rev)
 
 	def amino_acid_export(self, aa_transporters_counts: np.ndarray,
 			aa_conc: Union[units.Unum, np.ndarray], mechanistic_uptake: bool):
@@ -1349,20 +1338,12 @@ def amino_acid_export(self, aa_transporters_counts: np.ndarray,
 			rate of export for each amino acid. array is unitless but
 				represents counts of amino acid per second
 		"""
+		if units.hasUnit(aa_conc):
+			aa_conc = aa_conc.asNumber(METABOLITE_CONCENTRATION_UNITS)
 
-		if mechanistic_uptake:
-			# Export based on mechanistic model
-			if units.hasUnit(aa_conc):
-				aa_conc = aa_conc.asNumber(METABOLITE_CONCENTRATION_UNITS)
-			trans_counts_per_aa = self.aa_to_exporters_matrix @ aa_transporters_counts
-			export_rates = self.export_kcats_per_aa * trans_counts_per_aa / (1 + self.aa_export_kms / aa_conc)
-		else:
-			# Export is lumped with specific uptake rates in amino_acid_import
-			# and not dependent on internal amino acid concentrations or
-			# explicitly considered here
-			export_rates = np.zeros(len(aa_conc))
-
-		return export_rates
+		return amino_acid_export_jit(aa_transporters_counts, aa_conc,
+			mechanistic_uptake, self.aa_to_exporters_matrix,
+			self.export_kcats_per_aa, self.aa_export_kms)
 
 	def amino_acid_import(self, aa_in_media: np.ndarray, dry_mass: units.Unum,
 			internal_aa_conc: Union[units.Unum, np.ndarray], aa_transporters_counts: np.ndarray,
@@ -1385,15 +1366,11 @@ def amino_acid_import(self, aa_in_media: np.ndarray, dry_mass: units.Unum,
 		if units.hasUnit(internal_aa_conc):
 			internal_aa_conc = internal_aa_conc.asNumber(METABOLITE_CONCENTRATION_UNITS)
 
-		saturation = 1 / (1 + internal_aa_conc / self.aa_import_kis)
-		if mechanisitic_uptake:
-			# Uptake based on mechanistic model
-			counts_per_aa = self.aa_to_importers_matrix @ aa_transporters_counts
-			import_rates = self.import_kcats_per_aa * counts_per_aa
-		else:
-			import_rates = self.max_specific_import_rates * dry_mass.asNumber(DRY_MASS_UNITS)
-
-		return import_rates * saturation * aa_in_media
+		dry_mass = dry_mass.asNumber(DRY_MASS_UNITS)
+		return amino_acid_import_jit(aa_in_media, dry_mass, internal_aa_conc,
+			aa_transporters_counts, mechanisitic_uptake, self.aa_import_kis,
+			self.aa_to_importers_matrix, self.import_kcats_per_aa,
+			self.max_specific_import_rates)
 
 	def get_amino_acid_conc_conversion(self, conc_units):
 		return units.strip_empty_units(conc_units / METABOLITE_CONCENTRATION_UNITS)
@@ -2224,6 +2201,83 @@ def _is_transport_rxn(self, stoich):
 		return is_transport_rxn
 
 
+@njit
+def np_apply_along_axis(func1d, axis, arr):
+	if arr.ndim != 2:
+		raise RuntimeError("Array must have 2 dimensions.")
+	if axis not in [0, 1]:
+		raise RuntimeError("Axis must be 0 or 1.")
+	if axis == 0:
+		result = np.empty(arr.shape[1])
+		for i in range(len(result)):
+			result[i] = func1d(arr[:, i])
+	else:
+		result = np.empty(arr.shape[0])
+		for i in range(len(result)):
+			result[i] = func1d(arr[i, :])
+	return result
+
+
+@njit
+def np_prod(array, axis):
+	return np_apply_along_axis(np.prod, axis, array)
+
+
+@njit
+def amino_acid_synthesis_jit(counts_per_aa_fwd, counts_per_aa_rev, aa_conc,
+	aa_upstream_kms, aa_kis, aa_reverse_kms,
+	aa_degradation_kms, aa_forward_stoich, aa_kcats_fwd,
+	aa_reverse_stoich, aa_kcats_rev):
+	km_saturation = np_prod(1 / (1 + aa_upstream_kms / aa_conc), axis=1)
+
+	# Determine saturation fraction for reactions
+	forward_fraction = 1 / (1 + aa_conc / aa_kis) * km_saturation
+	reverse_fraction = 1 / (1 + aa_reverse_kms / aa_conc)
+	deg_fraction = 1 / (1 + aa_degradation_kms / aa_conc)
+	loss_fraction = reverse_fraction + deg_fraction
+
+	# Calculate synthesis rate
+	synthesis = (
+		aa_forward_stoich.astype(np.float64) @ (aa_kcats_fwd * counts_per_aa_fwd * forward_fraction).astype(np.float64)
+		- aa_reverse_stoich.astype(np.float64) @ (aa_kcats_rev * counts_per_aa_rev * reverse_fraction).astype(np.float64)
+		- aa_kcats_rev * counts_per_aa_rev * deg_fraction
+	)
+
+	return synthesis, forward_fraction, loss_fraction
+
+
+@njit
+def amino_acid_export_jit(aa_transporters_counts, aa_conc,
+	mechanistic_uptake, aa_to_exporters_matrix,
+	export_kcats_per_aa, aa_export_kms):
+	if mechanistic_uptake:
+		# Export based on mechanistic model
+		trans_counts_per_aa = aa_to_exporters_matrix.astype(np.float64) @ aa_transporters_counts.astype(np.float64)
+		export_rates = export_kcats_per_aa * trans_counts_per_aa / (1 + aa_export_kms / aa_conc)
+	else:
+		# Export is lumped with specific uptake rates in amino_acid_import
+		# and not dependent on internal amino acid concentrations or
+		# explicitly considered here
+		export_rates = np.zeros(len(aa_conc))
+	return export_rates
+
+
+@njit
+def amino_acid_import_jit(aa_in_media, dry_mass, internal_aa_conc,
+	aa_transporters_counts, mechanisitic_uptake, aa_import_kis,
+	aa_to_importers_matrix, import_kcats_per_aa,
+	max_specific_import_rates):
+	saturation = 1 / (1 + internal_aa_conc / aa_import_kis)
+	if mechanisitic_uptake:
+		# Uptake based on mechanistic model
+		counts_per_aa = aa_to_importers_matrix.astype(np.float64) @ aa_transporters_counts.astype(np.float64)
+		import_rates = import_kcats_per_aa * counts_per_aa
+	else:
+		import_rates = max_specific_import_rates * dry_mass
+
+	return import_rates * saturation * aa_in_media
+
+
 # Class used to update metabolite concentrations based on the current nutrient conditions
 class ConcentrationUpdates(object):
 	def __init__(self, concDict, relative_changes, equilibriumReactions, exchange_data_dict, all_metabolite_ids, linked_metabolites):

reconstruction/ecoli/dataclasses/process/transcription.py

@@ -5,6 +5,7 @@
 TODO: handle ppGpp and DksA-ppGpp regulation separately
 """
 
+from functools import cache
 from typing import cast
 
 import numpy as np
@@ -908,6 +909,7 @@ def cistron_id_to_rna_indexes(self, cistron_id):
 		return self.cistron_tu_mapping_matrix.getrow(
 			self._cistron_id_to_index[cistron_id]).nonzero()[1]
 
+	@cache
 	def rna_id_to_cistron_indexes(self, rna_id):
 		"""
 		Returns the indexes of cistrons that constitute the given transcription