Paper code: ED-AFM (#41)

.gitignore

@@ -81,5 +81,11 @@ dmypy.json
 # VS Code
 .vscode
 
+# Tex
+*.aux
+*.log
+*.pdf
+
 # Other
 molecules
+*.csv

README.md

@@ -30,4 +30,5 @@ pip install .
 ## Papers
 The [`papers`](papers) subdirectory contains training scripts and datasets for specific publications. Currently we have the following:
 - [Automated structure discovery in atomic force microscopy](papers/asd-afm)
+- [Electrostatic Discovery Atomic Force Microscopy](papers/ed-afm)
 - [Structure discovery in Atomic Force Microscopy imaging of ice](papers/ice_structure_discovery)

docs/source/reference/mlspm.models.rst

@@ -17,4 +17,12 @@ mlspm.models
 
     Alias of :class:`mlspm.image.models.ASDAFMNet`
 
+.. class:: mlspm.models.AttentionUNet
+
+    Alias of :class:`mlspm.image.models.AttentionUNet`
+
+.. class:: mlspm.models.EDAFMNet
+
+    Alias of :class:`mlspm.image.models.EDAFMNet`
+
 .. autofunction:: mlspm.models.download_weights

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,8 @@
     "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/10609676/files/molecules_rebias.tar.gz?download=1",
+    "ED-AFM-data": "https://zenodo.org/records/10609676/files/edafm-data.tar.gz?download=1",
 }
 
 
@@ -43,6 +45,8 @@ 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.10609676 - 'molecules_rebias.tar.gz'
+        - ``'ED-AFM-data'``: https://doi.org/10.5281/zenodo.10609676 - 'edafm-data.tar.gz'
 
     Arguments:
         name: Name of the dataset to download.