Moved RML imaging inside the snakemake workflow.

02-sgd/README.md

@@ -26,27 +26,23 @@ snakemake -c1 all
 
 # Description of Contents
 
-This repository assumes that you will run all scripts from this `02` directory (the one containing `02-sgd/README.md`). Some aspects of the workflow are automated with Snakemake ([`Snakefile`](Snakefile)). 
+This repository assumes that you will run all scripts from this `02` directory (the one containing `02-sgd/README.md`). All aspects of a basic workflow are automated as above with Snakemake ([`Snakefile`](Snakefile)), but we recommend looking at the individual Snakemake rules and the associated Python scripts to understand what is happening. 
 
 First, we recommend looking at [`src/load_data.py`](src/load_data.py) to see how mock visibilities $\mathcal{V}(u,v)$ are generated from the mock image and baselines.
 
 Then, we recommend looking at [`src/plot_baselines.py`](src/plot_baselines.py) and [`src/dirty_image.py`](src/dirty_image.py) to make diagnostic plots of the baseline and a dirty image of the data, to check that everything appears as you might expect.
 
-You can run these simple scripts using 
-
-```
-$ snakemake -c1 all
-```
+These scripts produce the following plots
 
 ![baselines](analysis/baselines.png)
 
 ![Dirty Beam and Image](analysis/dirty_image.png)
 
 # RML imaging workflow
 
-The RML imaging workflow is demonstrated in [`src/sgd.py`](src/sgd.py). We recommend looking through that file before reading the rest of this document. If you are new to PyTorch idioms, we recommend familiarizing yourself with the [PyTorch basics](https://mpol-dev.github.io/MPoL/background.html#pytorch) first. 
+The RML imaging workflow is demonstrated in [`src/sgd.py`](src/sgd.py). We recommend looking through that file before reading the rest of this document. If you are new to PyTorch idioms, we recommend familiarizing yourself with the [PyTorch basics](https://mpol-dev.github.io/MPoL/background.html#pytorch) first.
 
-The RML imaging workflow is not part of the Snakemake workflow, instead, one runs the script like
+One can run the RML imaging workflow outside of the Snakemake workflow by, for example,
 
 ```
 $ python src/sgd.py --epochs=5

02-sgd/Snakefile

@@ -1,7 +1,8 @@
 rule all:
     input: 
         "analysis/baselines.png",
-        "analysis/dirty_image.png"        
+        "analysis/dirty_image.png",
+        "analysis/butterfly.png"
 
 rule baselines:
     input: "data/mock_data.npz", "src/plot_baselines.py"
@@ -14,6 +15,16 @@ rule dirty_image:
     shell: "python src/dirty_image.py {output}"
 
 
+rule checkpoint_en0:
+    input: "data/mock_data.npz", "src/sgd.py"
+    output: "checkpoints/ent0.pt"
+    shell: "python src/sgd.py --tensorboard-log-dir=runs/ent0 --save-checkpoint=checkpoints/ent0.pt --lr 1e-1 --FWHM 0.05 --epochs=5 --lam-ent=1e-5"
+
+rule rml_image:
+    input: checkpoint="checkpoints/ent0.pt", src="src/plot_image.py"
+    output: "analysis/butterfly.png"
+    shell: "python {input.src} {input.checkpoint} {output}"
+
 # some initial attempts
 # python src/sgd.py --tensorboard-log-dir=runs/exp0 --save-checkpoint=checkpoints/0.pt --lr 1e-2 --FWHM 0.05 --epochs=50
 # python src/sgd.py --tensorboard-log-dir=runs/exp1 --load-checkpoint=checkpoints/0.pt --save-checkpoint=checkpoints/1.pt --lr 1e-1 --FWHM 0.05 --epochs=30