Added links to ED-AFM paper pretrained weights and molecule dataset.

Author: NikoOinonen

mlspm/_weights.py

@@ -11,6 +11,14 @@
     "graph-ice-au111-bilayer": "https://zenodo.org/records/10054348/files/weights_ice-au111-bilayer.pth?download=1",
     "asdafm-light": "https://zenodo.org/records/10514470/files/weights_asdafm_light.pth?download=1",
     "asdafm-heavy": "https://zenodo.org/records/10514470/files/weights_asdafm_heavy.pth?download=1",
+    "edafm-base": "https://zenodo.org/records/10606273/files/base.pth?download=1",
+    "edafm-single-channel": "https://zenodo.org/records/10606273/files/single-channel.pth?download=1",
+    "edafm-CO-Cl": "https://zenodo.org/records/10606273/files/CO-Cl.pth?download=1",
+    "edafm-Xe-Cl": "https://zenodo.org/records/10606273/files/Xe-Cl.pth?download=1",
+    "edafm-constant-noise": "https://zenodo.org/records/10606273/files/constant-noise.pth?download=1",
+    "edafm-uniform-noise": "https://zenodo.org/records/10606273/files/uniform_noise.pth?download=1",
+    "edafm-no-gradient": "https://zenodo.org/records/10606273/files/no-gradient.pth?download=1",
+    "edafm-matched-tips": "https://zenodo.org/records/10606273/files/matched-tips.pth?download=1",
 }
 
 
@@ -21,11 +29,30 @@ def download_weights(weights_name: str, target_path: Optional[PathLike] = None)
     The following weights are available:
 
         - ``'graph-ice-cu111'``: PosNet trained on ice clusters on Cu(111). (https://doi.org/10.5281/zenodo.10054348)
-        - ``'graph-ice-au111-monolayer'``: PosNet trained on monolayer ice clusters on Au(111). (https://doi.org/10.5281/zenodo.10054348)
-        - ``'graph-ice-au111-bilayer'``: PosNet trained on bilayer ice clusters on Au(111). (https://doi.org/10.5281/zenodo.10054348)
-        - ``'asdafm-light'``: ASDAFMNet trained on molecules containing the elements H, C, N, O, and F. (https://doi.org/10.5281/zenodo.10514470)
-        - ``'asdafm-heavy'``: ASDAFMNet trained on molecules additionally containing Si, P, S, Cl, and Br. (https://doi.org/10.5281/zenodo.10514470)
-
+        - ``'graph-ice-au111-monolayer'``: PosNet trained on monolayer ice clusters on Au(111).
+          (https://doi.org/10.5281/zenodo.10054348)
+        - ``'graph-ice-au111-bilayer'``: PosNet trained on bilayer ice clusters on Au(111).
+          (https://doi.org/10.5281/zenodo.10054348)
+        - ``'asdafm-light'``: :class:`.ASDAFMNet` trained on molecules containing the elements H, C, N, O, and F.
+          (https://doi.org/10.5281/zenodo.10514470)
+        - ``'asdafm-heavy'``: :class:`.ASDAFMNet` trained on molecules additionally containing Si, P, S, Cl, and Br.
+          (https://doi.org/10.5281/zenodo.10514470)
+        - ``'edafm-base'``: :class:`.EDAFMNet` used for all predictions in the main ED-AFM paper and used for comparison in
+          the various tests in the supplementary information of the paper. (https://doi.org/10.5281/zenodo.10606273)
+        - ``'edafm-single-channel'``: :class:`.EDAFMNet` trained on only a single CO-tip AFM input.
+          (https://doi.org/10.5281/zenodo.10606273)
+        - ``'edafm-CO-Cl'``: :class:`.EDAFMNet` trained on alternative tip combination of CO and Cl.
+          (https://doi.org/10.5281/zenodo.10606273)
+        - ``'edafm-Xe-Cl'``: :class:`.EDAFMNet` trained on alternative tip combination of Xe and Cl.
+          (https://doi.org/10.5281/zenodo.10606273)
+        - ``'edafm-constant-noise'``: :class:`.EDAFMNet` trained using constant noise amplitude instead of normally distributed
+          amplitude. (https://doi.org/10.5281/zenodo.10606273)
+        - ``'edafm-uniform-noise'``: :class:`.EDAFMNet` trained using uniform random noise amplitude instead of normally
+          distributed amplitude. (https://doi.org/10.5281/zenodo.10606273)
+        - ``'edafm-no-gradient'``: :class:`.EDAFMNet` trained without background-gradient augmentation.
+          (https://doi.org/10.5281/zenodo.10606273)
+        - ``'edafm-matched-tips'``: :class:`.EDAFMNet` trained on data with matched tip distance between CO and Xe,
+          instead of independently randomized distances. (https://doi.org/10.5281/zenodo.10606273)
 
     Arguments:
         weights_name: Name of weights to download.

mlspm/datasets.py

@@ -15,6 +15,7 @@
     "AFM-ice-relaxed": "https://zenodo.org/records/10362511/files/relaxed_structures.tar.gz?download=1",
     "ASD-AFM-molecules": "https://zenodo.org/records/10562769/files/molecules.tar.gz?download=1",
     "AFM-camphor-exp": "https://zenodo.org/records/10562769/files/afm_camphor.tar.gz?download=1",
+    "ED-AFM-molecules": "https://zenodo.org/records/10606443/files/molecules_rebias.tar.gz?download=1",
 }
 
 
@@ -43,6 +44,7 @@ def download_dataset(name: str, target_dir: PathLike):
         - ``'AFM-ice-relaxed'``: https://doi.org/10.5281/zenodo.10362511
         - ``'ASD-AFM-molecules'``: https://doi.org/10.5281/zenodo.10562769 - 'molecules.tar.gz'
         - ``'AFM-camphor-exp'``: https://doi.org/10.5281/zenodo.10562769 - 'afm_camphor.tar.gz'
+        - ``'ED-AFM-molecules'``: https://doi.org/10.5281/zenodo.10606443
 
     Arguments:
         name: Name of the dataset to download.

mlspm/image/models.py

@@ -396,7 +396,6 @@ class EDAFMNet(AttentionUNet):
     Arguments:
         device: Device to load model onto.
         trained_weights: If not None, load the specified pretrained weights to the model.
-        weights_dir: If **weights_type** is not None, the directory where the weights will be downloaded into.
     """
 
     def __init__(
@@ -405,7 +404,6 @@ def __init__(
         pretrained_weights: Optional[
             Literal["base", "single-channel", "CO-Cl", "Xe-Cl", "constant-noise", "uniform-noise", "no-gradient", "matched-tips"]
         ] = None,
-        weights_dir: PathLike = "./weights",
     ):
         if pretrained_weights == "single-channel":
             n_in = 1
@@ -440,7 +438,7 @@ def __init__(
         )
 
         if pretrained_weights:
-            weights_path = download_weights(f"ed-afm-{pretrained_weights}", weights_dir)
+            weights_path = download_weights(f"edafm-{pretrained_weights}")
             self.load_state_dict(torch.load(weights_path))
 
 

papers/ed-afm/README.md

@@ -28,11 +28,11 @@ In `EDAFMNet` one can also specify pretrained weights of several types to downlo
  - 'no-gradient': Model trained without background-gradient augmentation.
  - 'matched-tips': Model trained on data with matched tip distance between CO and Xe, instead of independently randomized distances.
 
-The model weights can also be downloaded directly from https://www.dropbox.com/sh/1jd4h7ovpafg3uk/AACQX6C6Nks6RxGbvqb6vIo-a?dl=0. The weights are saved in the state_dict format of PyTorch.
+The model weights can also be downloaded directly from https://doi.org/10.5281/zenodo.10606273. The weights are saved in the state_dict format of PyTorch.
 
 ## Data and model training
 
-We don't provide the full training/validation/test sets for download because they are very large (~1TiB in total). Instead, we provide the database of molecular geometries that can be used to generate the full dataset using ProbeParticleModel. The provided script `generate_data.py` does the data generation and will download the molecule database automatically. Alternatively, the molecule database can be downloaded directly from https://www.dropbox.com/s/z4113upq82puzht/Molecules_rebias_210611.tar.gz?dl=0. 
+We don't provide the full training/validation/test sets for download because they are very large (~1TiB in total). Instead, we provide the database of molecular geometries that can be used to generate the full dataset using ProbeParticleModel. The provided script `generate_data.py` does the data generation and will download the molecule database automatically. Alternatively, the molecule database can be downloaded directly from https://doi.org/10.5281/zenodo.10606443.
 
 The model training can be done using the provided script `train.py`. Note that performing the training using all the same settings as we used requires a significant amount of time and also a significant amount VRAM on the GPU, likely more than can be found on a single GPU. In our case the model training took ~5 days using 4 x Nvidia Tesla V100 (32GB) GPUs. However, inference on the trained model can be done even on a single lower-end GPU or on CPU.