Skip to content

Commit

Permalink
Refactor code to improve performance and readability
Browse files Browse the repository at this point in the history
  • Loading branch information
@sulfierry
sulfierry committed Mar 21, 2024
1 parent 583ecb3 commit 8b1f8cd
Showing 1 changed file with 63 additions and 19 deletions.
82 changes: 63 additions & 19 deletions free_energy_landscape/freeEnergyLandscape.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,9 @@ class FreeEnergyLandscape:
def __init__(self, cv1_path, cv2_path,
temperature, boltzmann_constant,
bins=100, kde_bandwidth=None,
cv_names=['CV1', 'CV2'], discrete=None):
cv_names=['CV1', 'CV2'], discrete=None,
xlim_inf=None, xlim_sup=None, ylim_inf=None, ylim_sup=None):


self.cv1_path = cv1_path
self.cv2_path = cv2_path
Expand Down Expand Up @@ -51,7 +53,7 @@ def __init__(self, cv1_path, cv2_path,
'purple',
'darkorange',
'green',
'orange',
'lightgrey',
'red',
'magenta',
'mediumorchid',
Expand All @@ -73,6 +75,12 @@ def __init__(self, cv1_path, cv2_path,
'd'
]

# Novos atributos para limites dos eixos
self.xlim_inf = xlim_inf
self.xlim_sup = xlim_sup
self.ylim_inf = ylim_inf
self.ylim_sup = ylim_sup

def load_data(self):
# Carrega os dados das variáveis coletivas e os índices dos frames
self.proj1_data_original = np.loadtxt(self.cv1_path, usecols=[1])
Expand All @@ -88,19 +96,26 @@ def calculate_free_energy(self, data):
kernel_original = gaussian_kde(values_original.T, bw_method=self.kde_bandwidth)
else:
kernel_original = gaussian_kde(values_original.T)
X_original, Y_original = np.mgrid[data[:, 0].min():data[:, 0].max():100j,
data[:, 1].min():data[:, 1].max():100j]

# Ajusta a geração da grade para respeitar os limites especificados
x_min = self.xlim_inf if self.xlim_inf is not None else data[:, 0].min()
x_max = self.xlim_sup if self.xlim_sup is not None else data[:, 0].max()
y_min = self.ylim_inf if self.ylim_inf is not None else data[:, 1].min()
y_max = self.ylim_sup if self.ylim_sup is not None else data[:, 1].max()

X_original, Y_original = np.mgrid[x_min:x_max:100j, y_min:y_max:100j]
positions_original = np.vstack([X_original.ravel(), Y_original.ravel()])
Z_original = np.reshape(kernel_original(positions_original).T, X_original.shape)
G_original = -self.kB * self.temperature * np.log(Z_original)
G_original = np.clip(G_original - np.min(G_original), 0, 25)

self.cached_results = {'X_original': X_original,
'Y_original': Y_original,
'G_original': G_original
}
'Y_original': Y_original,
'G_original': G_original
}
return self.cached_results


def boltzmann_inversion(self, data_list, titles, threshold=None):

fig_combined, axs_combined = plt.subplots(1, len(data_list), figsize=(20, 6), sharey=True)
Expand Down Expand Up @@ -230,7 +245,13 @@ def cv_by_frame(self, data_list, titles):
plt.show()


def plot_energy_landscape(self, threshold, titles=['CV1', 'CV2']):
def plot_energy_landscape(self, threshold, titles=['CV1', 'CV2'], xlim_inf=None, xlim_sup=None, ylim_inf=None, ylim_sup=None):

if xlim_inf is not None and xlim_sup is not None:
plt.xlim(xlim_inf, xlim_sup)
if ylim_inf is not None and ylim_sup is not None:
plt.ylim(ylim_inf, ylim_sup)

data = np.hstack((self.proj1_data_original[:, None], self.proj2_data_original[:, None]))
result = self.calculate_free_energy(data)
plt.figure(figsize=(8, 6))
Expand Down Expand Up @@ -270,14 +291,16 @@ def plot_energy_landscape(self, threshold, titles=['CV1', 'CV2']):

cbar = plt.colorbar(cont)
cbar.set_label('Free energy (kJ/mol)')
plt.xlim(self.xlim_inf, self.xlim_sup)
plt.ylim(self.ylim_inf, self.ylim_sup)
plt.xlabel(titles[0])
plt.ylabel(titles[1])
plt.title('Free Energy Landscape')
plt.tight_layout(rect=[0, 0, 0.85, 1]) # Ajuste para garantir que a legenda fique visível e não sobreponha o gráfico
plt.savefig('Free_energy_landscape_with_discrete_points.png')
plt.show()

def plot_3D_energy_landscape(self, threshold=None, titles=['CV1', 'CV2']):
def plot_3D_energy_landscape(self, threshold=None, titles=['CV1', 'CV2'], xlim_inf=None, xlim_sup=None, ylim_inf=None, ylim_sup=None):
data = np.hstack((self.proj1_data_original[:, None], self.proj2_data_original[:, None]))
result = self.calculate_free_energy(data)

Expand Down Expand Up @@ -321,7 +344,9 @@ def plot_3D_energy_landscape(self, threshold=None, titles=['CV1', 'CV2']):
)

ax.legend(loc='upper left', bbox_to_anchor=(1.05, 1), title="Energy intervals")


plt.xlim(self.xlim_inf, self.xlim_sup)
plt.ylim(self.ylim_inf, self.ylim_sup)
plt.tight_layout()
plt.show()

Expand All @@ -335,10 +360,7 @@ def plot_threshold_points(self, ax, result, lower_bound, upper_bound, color, lab
ax.scatter(X_flat[energy_mask], Y_flat[energy_mask], G_flat[energy_mask], color=color, s=20, label=label)


def create_3D_gif(self, gif_filename='energy_landscape_3D.gif',
n_angles=10, elevation=15,
duration_per_frame=0.01,
titles=['CV1', 'CV2']):
def create_3D_gif(self, gif_filename='energy_landscape_3D.gif', n_angles=10, elevation=15, duration_per_frame=0.01, titles=['CV1', 'CV2'], xlim_inf=None, xlim_sup=None, ylim_inf=None, ylim_sup=None):

temp_dir = tempfile.mkdtemp() # Cria um diretório temporário para armazenar os frames
filenames = []
Expand Down Expand Up @@ -415,9 +437,13 @@ def help():
--gif_angles [int] Angles for 3D GIF rotation (default: 10)
--gif_elevation [int] Elevation angle for the 3D GIF (default: 10)
--gif_duration [float] Duration per frame in the GIF in seconds (default: 0.1)
--xlim_inf [float] Lower limit for the x-axis (default: None)
--xlim_sup [float] Upper limit for the x-axis (default: None)
--ylim_inf [float] Lower limit for the y-axis (default: None)
--ylim_sup [float] Upper limit for the y-axis (default: None)
Example:
free_energy_landscape cv1.txt cv2.txt --names "Angle (CV1)" "Distance (CV2)" --energy 3.0 --discretize 1.0
freeEnergyLandscape.py proj1Out.txt proj2Out.txt --energy 0.1 --discretize 0.1 --xlim_inf 27 --xlim_sup 130 --ylim_inf 3 --ylim_sup 37
"""
print(help_text)
Expand Down Expand Up @@ -547,6 +573,8 @@ def main():
elevation = 10 # --gif_elevation [int]
duration_per_frame = 0.1 # --gif_duration [float]
discrete_val = None # --discrete [float]
xlim_inf = xlim_sup = ylim_inf = ylim_sup = None # Inicialização padrão


if len(sys.argv) >= 3:
cv1_path, cv2_path = sys.argv[1], sys.argv[2]
Expand Down Expand Up @@ -585,6 +613,20 @@ def main():
elif key == "--gif_duration":
duration_per_frame = float(sys.argv[i + 1])
i += 2

elif key == "--xlim_inf":
xlim_inf = float(sys.argv[i + 1])
i += 2
elif key == "--xlim_sup":
xlim_sup = float(sys.argv[i + 1])
i += 2
elif key == "--ylim_inf":
ylim_inf = float(sys.argv[i + 1])
i += 2
elif key == "--ylim_sup":
ylim_sup = float(sys.argv[i + 1])
i += 2

else:
print(f"Unrecognized option: {key}")
sys.exit(1)
Expand All @@ -594,10 +636,12 @@ def main():

try:
fel = FreeEnergyLandscape(cv1_path, cv2_path, t, kB,
bins=bins_energy_histogram,
kde_bandwidth=kde_bandwidth_cv,
cv_names=cv_names,
discrete=discrete_val)
bins=bins_energy_histogram,
kde_bandwidth=kde_bandwidth_cv,
cv_names=cv_names,
discrete=discrete_val,
xlim_inf=xlim_inf, xlim_sup=xlim_sup,
ylim_inf=ylim_inf, ylim_sup=ylim_sup)

fel.main(energy_threshold, cv_names=cv_names,
n_angles=n_angles, elevation=elevation,
Expand Down

0 comments on commit 8b1f8cd

Please sign in to comment.