From d42bf153485007c4167a09e52748d64f8d2b293c Mon Sep 17 00:00:00 2001 From: DSilva27 Date: Wed, 23 Aug 2023 14:42:10 -0400 Subject: [PATCH] implement solvent model --- src/cryo_sbi/inference/priors.py | 14 +++++----- .../wpa_simulator/image_generation.py | 27 +++++++++++++++++++ 2 files changed, 34 insertions(+), 7 deletions(-) diff --git a/src/cryo_sbi/inference/priors.py b/src/cryo_sbi/inference/priors.py index be15a14..9608f0d 100644 --- a/src/cryo_sbi/inference/priors.py +++ b/src/cryo_sbi/inference/priors.py @@ -44,10 +44,10 @@ def get_image_priors( ) assert ( - lower > 2.0 * image_config["PIXEL_SIZE"] + lower >= 2.0 * image_config["PIXEL_SIZE"] ), "The lower bound for RES must be at least 2 times the pixel size." - assert lower < upper, "Lower bound must be smaller than upper bound." + assert lower < upper, "Lower bound for RES must be smaller than upper bound." res = zuko.distributions.BoxUniform(lower=lower, upper=upper, ndims=1) @@ -74,7 +74,7 @@ def get_image_priors( ) assert lower > 0.0, "The lower bound for DEFOCUS must be positive." - assert lower < upper, "Lower bound must be smaller than upper bound." + assert lower < upper, "Lower bound for DEFOCUS must be smaller than upper bound." defocus = zuko.distributions.BoxUniform(lower=lower, upper=upper, ndims=1) @@ -89,8 +89,8 @@ def get_image_priors( [[image_config["B_FACTOR"][1]]], dtype=torch.float32, device=device ) - assert lower > 0.0, "The lower bound for DEFOCUS must be positive." - assert lower < upper, "Lower bound must be smaller than upper bound." + assert lower > 0.0, "The lower bound for B_FACTOR must be positive." + assert lower < upper, "Lower bound for B_FACTOR must be smaller than upper bound." b_factor = zuko.distributions.BoxUniform(lower=lower, upper=upper, ndims=1) @@ -102,8 +102,8 @@ def get_image_priors( [[image_config["SNR"][1]]], dtype=torch.float32, device=device ).log10() - assert lower > 0.0, "The lower bound for DEFOCUS must be positive." - assert lower < upper, "Lower bound must be smaller than upper bound." + #assert lower > 0.0, "The lower bound for DEFOCUS must be positive." + #assert lower < upper, "Lower bound must be smaller than upper bound." snr = zuko.distributions.BoxUniform(lower=lower, upper=upper, ndims=1) diff --git a/src/cryo_sbi/wpa_simulator/image_generation.py b/src/cryo_sbi/wpa_simulator/image_generation.py index 614d634..16204f8 100644 --- a/src/cryo_sbi/wpa_simulator/image_generation.py +++ b/src/cryo_sbi/wpa_simulator/image_generation.py @@ -89,6 +89,7 @@ def project_density( coords_rot = torch.bmm(rot_matrix, atomic_model[:, :3, :]) coords_rot[:, :2, :] += shift.unsqueeze(-1) + # Protein Projection gauss_x = torch.exp_( -((grid.unsqueeze(-1) - coords_rot[:, 0, :].unsqueeze(1)) ** 2) / variances.unsqueeze(1) @@ -101,4 +102,30 @@ def project_density( image = torch.bmm(gauss_x, gauss_y.transpose(1, 2)) + # Implicit Solvent Projection + freq_cutoff = 1.0 / 6.0 # 6 angstrom + solv_coeff = 4.0 # standard deviation is twice the one for atoms + + variances = solv_coeff * atomic_model[:, 4, :] * res[:, 0] ** 2 + amplitudes = atomic_model[:, 3, :] / torch.sqrt((2 * torch.pi * variances)) + + gauss_x = torch.exp_( + -((grid.unsqueeze(-1) - coords_rot[:, 0, :].unsqueeze(1)) ** 2) + / variances.unsqueeze(1) + ) * amplitudes.unsqueeze(1) + + gauss_y = torch.exp( + -((grid.unsqueeze(-1) - coords_rot[:, 1, :].unsqueeze(1)) ** 2) + / variances.unsqueeze(1) + ) * amplitudes.unsqueeze(1) + + image_solv = torch.bmm(gauss_x, gauss_y.transpose(1, 2)) + + freq_pix_1d = torch.fft.fftfreq(int(num_pixels), d=pixel_size, device=image_solv.device) + freq2_2d = freq_pix_1d[:, None] ** 2 + freq_pix_1d[None, :] ** 2 + + image_solv = 0.5 * torch.fft.ifft2(torch.fft.fft2(image_solv) * (freq2_2d < freq_cutoff)).real + + image = image - image_solv + return image