implement solvent model

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)
 

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