working example

.devcontainer/devcontainer.json

@@ -21,7 +21,7 @@
 	// "features": {},
 
 	// Use 'forwardPorts' to make a list of ports inside the container available locally.
-	"forwardPorts": [7860],
+	// "forwardPorts": [],
 
 	// Uncomment the next line to run commands after the container is created.
 	// "postCreateCommand": "cat /etc/os-release",

.gitattributes

@@ -37,3 +37,17 @@ data/training_data.pkl filter=lfs diff=lfs merge=lfs -text
 *.pkl filter=lfs diff=lfs merge=lfs -text
 **/*.pkl filter=lfs diff=lfs merge=lfs -text
 data filter=lfs diff=lfs merge=lfs -text
+data/weights/h1_w2vec/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/h7_w2vec/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/r3_local/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/r3_w2vec/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/r5_w2vec/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/r9_local/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/r9_nbrs/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/r9_w2vec/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/h3_w2vec/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/r7_w2vec/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/r1_w2vec/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/r3_nbrs/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/h5_w2vec/model.pt filter=lfs diff=lfs merge=lfs -text
+data/weights/h9_w2vec/model.pt filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,4 +1,5 @@
-from shiny import App, ui, render
+from shiny import App, ui, render, reactive
+from shiny.ui import HTML, tags
 import shinyswatch
 
 import torch
@@ -12,10 +13,9 @@ from collections import defaultdict
 from tqdm import tqdm
 import itertools as it
 from torch import nn
+import io
 
-import pickle
-from models import UNetEncoder, Decoder
-from utils import load_training_data
+from utils import load_training_data, load_models
 
 MONTHS= {
     0: "Jan",
@@ -61,35 +61,139 @@ C, NAMES, Y, M = load_training_data(
 )
 _, _, YRAW, MRAW  = load_training_data(path="data/training_data.pkl")
 
+prefix = "h"
+dirs = {
+    "r1": f"./data/weights/{prefix}1_w2vec",
+    "r3": f"./data/weights/{prefix}3_w2vec",
+    "r5": f"./data/weights/{prefix}5_w2vec",
+    "r7": f"./data/weights/{prefix}7_w2vec",
+    "r9": f"./data/weights/{prefix}9_w2vec",
+}
+MODELS = load_models(dirs, prefix=prefix, nd=5)
+
+
+multicol_html = tags.head(
+  tags.style(
+    HTML(
+        ".multicol {"
+        # "height: 150px; "
+        "-webkit-column-count: 3;" # chrome, safari, opera
+        "-moz-column-count: 3;"  # firefox
+        "column-count: 3;"
+        "-moz-column-fill: auto;"
+        "-column-fill: auto;"
+    )
+  )
+)
+
+instructions = f"""
+### Instructions
+
+Upload a CSV file with columns (id, lat, lon) using the `Browse` button on the sidebar.
+Below is an example of the contents of the file:
+
+
+```
+id,lat,lon
+0,47.5,-122.5
+1,47.5,-122.25
+2,47.5,-122.0
+3,47.5,-121.75
+4,47.5,-121.5
+```
+
+
+The id column can be any identifier, or the column can be ommited, in which case the row number will be used as the id.
+Make sure that the latitude is before the longitude column in the CSV file. The valid range for latitude is
+{YMIN} to {YMAX} and longitude is {XMIN} to {XMAX}, which cover the contiguous United States. 
+
+The resolution corresponds to how much neighboring information is captured by the embedding. If `local` is selected,
+the original weather covariates will be returned. Currently, all the embeddings correspond to the variables:
+air temperature (2m), precipitation, relative humidity (2m), vertical wind speed (10m), and horizontal wind speed (10m).
+The native resolution of the covariates is ~32 km for a grid size of 128 x 256. 
+
+### Results
+
+"""
+
+# After uploading the file, the app will generate a CSV, a download link will appear here.
+# The CSV will contain the following columns:
+
 
 # Part 1: ui ----
 app_ui = ui.page_fluid(
     shinyswatch.theme.minty(),
-    ui.panel_title("Generate Weather2vec Embeddings"),
-    ui.panel_sidebar(
-        ui.input_file("df", "Upload CSV File", accept=".csv"),
-        ui.input_checkbox_group("months", "Select Months", MONTHS),
-        ui.input_checkbox_group("months", "Select Years", YEARS),
-        ui.input_checkbox_group("months", "Select Resolutions", RESOLUTIONS),
-        width=6,
-    ),
-    ui.panel_main(
-        ui.output_plot("plot"),
-    ),
+    multicol_html,
+    ui.panel_title("Welcome to the Weather2vec Embedding Generator!"),
+    ui.layout_sidebar(
+        ui.panel_sidebar(
+            ui.input_file("df", "Upload CSV File", accept=".csv"),
+            tags.div(
+                ui.input_checkbox_group("months", HTML("<b>Months</b>"), MONTHS),
+                class_="multicol",
+                align="left",
+                inline=False
+            ),
+            HTML("<b>Note:</b> Embedding of multiple months will be added.<br>True multi-temporal embeddings will be supported in the future.<br><br>"),
+            tags.div(
+                ui.input_radio_buttons("years", HTML("<b>Year</b>"), YEARS),
+                class_="multicol",
+                align="left",
+                inline=False
+            ),
+            HTML("<br>"),
+            tags.div(
+                ui.input_radio_buttons("resolutions", HTML("<b>Resolution</b>"), RESOLUTIONS),
+                class_="multicol",
+                align="left",
+                inline=False
+            ),
+            width=3,
+        ),
+        ui.panel_main(
+            ui.markdown(instructions),
+            ui.download_button("download", "Download Embeddings"),
+        ),
+    )
 )
 
 
 # Part 2: server ----
 def server(input, output, session):
     # make a plot
-    @output
-    @render.plot
-    def plot():
-        fig, ax = plt.subplots()
-        ax.plot(C[0, 0, 0])
-        return fig 
+    @session.download(filename="embeddings.csv")
+    def download():
+        # read input file
+        print(input.df()[-1].keys())
+        fname = input.df()[-1]['datapath']
+        df = pd.read_csv(fname)
+
+        # dfcols = []
+        # for k, v in D.items():
+        #     mod = v["mod"]
+        #     mod.load_state_dict(torch.load(os.path.join(k, "model.pt")))
+        #     with torch.no_grad():
+        #         Z = mod["enc"](Ct)
+        #         Z = Z.mean(0).cpu().numpy()
+        #         Zmat = Z[:, row, col].T
+        #         colnames = [f"C{i:02d}" for i in range(Zmat.shape[-1])]
+        #         Z = pd.DataFrame(Zmat, columns=colnames)
+        #         Z = pd.DataFrame(Zmat, columns=[x + f"_{len(dfcols)}" for x in colnames])
+        #         dfcols.append(Z)
+        #         Z = pd.concat([locs, Z], axis=1)
+        #         Z.to_csv(f"{savedir}/kms_{32 * int(v['radius']):03d}.csv", index=False)
+
+        # Cloc = Corig[ix].mean(0)[:, row, col].T
+        # Z = pd.DataFrame(Cloc, columns=colnames[:Cloc.shape[1]])
+        # Z = pd.concat([locs, Z], axis=1)
+        # Z.to_csv(f"{savedir}/kms_000.csv", index=False)   
 
+        with io.BytesIO() as f:
+            df.to_csv(f, index=False)
+            yield f.getvalue()
 
+    # # dfcols = pd.concat(dfcols, axis=1)
+    # # dfcols = pd.concat([Z, dfcols], axis=1)
 
 # Combine into a shiny app.
 # Note that the variable must be "app".

data/test-data.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:053c897211e68701fc2930f66e30c30d319c9858ae8f89a5293999a86da1c1c8
+size 35225

data/weights/h1_w2vec/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1350b39d48c739af0bb100b777c5ecab4e2ba5d2da9ec0020988d265aee9fecd
+size 320

data/weights/h1_w2vec/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a9a3f67afdfe9ed2f2f08b6e71872c215865b4c61405115a79c7e26c9688111
+size 492319

data/weights/h1_w2vec/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6898b26aafd53d108e078214fed89eae06caa5294bbda98b4b931d5d8cdffafc
+size 123

data/weights/h3_w2vec/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6852ae1a0b4514ee2391bb10e4e57563faeb252b447c672390dccb0551f40167
+size 320

data/weights/h3_w2vec/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a81e01e7ce2048ebb19f9e8446cf626aefacf334539687ff8f488ca9483e98e
+size 492319

data/weights/h3_w2vec/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24991bdccad89102340b4f77f062dd1db046055efacdba6cecc507c2612525f4
+size 22518

data/weights/h5_w2vec/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ea17a68d758e186172b873db76bf18c27eccbb32fb60b21b112603b75cd9cca0
+size 320

data/weights/h5_w2vec/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88611910163b4a9ddb7324a688c6b964321d821629250d38ef12624f4204783e
+size 492319

data/weights/h5_w2vec/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:525e8d69818e41e807220e6eead955704a31353dc6153ed8b964eab4e7e51f5b
+size 22519

data/weights/h7_w2vec/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c92d6f6fdd8f0e4fa19bfb7257e2fd3b3319bd418895d79dbfd50ff1365d13f1
+size 320

data/weights/h7_w2vec/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:150fdad043c8a9bf3c4548b97b019216c02dd241f4a267d544af1f6464b0856a
+size 492319

data/weights/h7_w2vec/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37f45005ff95ea52753df4b6771fdaae9c1320d84817b89e4d835d5f1ea7a70e
+size 22520

data/weights/h9_w2vec/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24098fc0d0c99e75b09345d27af6793b438383010d8dc56fabea5857769f8cea
+size 320

data/weights/h9_w2vec/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e39a8c1a079f68c7bb4ef90aa7cdd5b797474c582670f730d4d371e389c0ebf
+size 492319

data/weights/h9_w2vec/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6037749e8582a3954d0701ee8212497f8e2c75ab6764ef9f2e20614a306d2bcd
+size 22522

data/weights/r1_w2vec/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f438e42212b15d6eff79ef4fef2fa4e8a7af19567b24dda0ff7189dccddb4f83
+size 298

data/weights/r1_w2vec/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:77c8865887c1eec2a710c7aa8f54977564b0e76276e427497b64d6122926ced9
+size 4866847

data/weights/r1_w2vec/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7da97ac5558d1c2d220fed659d64b6e8c83041ee53e129802a1aa17ba8dd6a4d
+size 22003

data/weights/r3_local/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb68f0f9f23c5b0e9465357a032516f3d82da4d1e654db9e7a69fe311c802dd9
+size 320

data/weights/r3_local/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63e59b43d2633c39cb0b0906d052291b918a6c57d879452c956e1b13cd50c824
+size 35323

data/weights/r3_local/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bfe90fc358e1068101c5a089be7c7ea6c7716e1fdcaefa221906185244f9b87d
+size 21634

data/weights/r3_nbrs/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b121a5683b604ed17120febdf94f66f628fcf5e1df818ef3a12f5712bb402751
+size 320

data/weights/r3_nbrs/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:015b4cbe386bd17782869275d5ff1db5ecc9b45a685600a81e617bd9795861ce
+size 35323

data/weights/r3_nbrs/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24e28b6418688b93c2e10854ee2f95d66292efd99a19a01c0b2ab0f16d22d43b
+size 21633

data/weights/r3_w2vec/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d3471f6a8bbb72c9eac7f674601969edf9ae5830bf75556604215ad11f8905a
+size 298

data/weights/r3_w2vec/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1dd70f70e737d7ad349f0cde7661d509bd593f146a6de08710ef585376d37f13
+size 4866847

data/weights/r3_w2vec/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b73925d08e1f71f7333765fa5109b77e41e31f6ac2d35ff0b4351c83145bc321
+size 22139

data/weights/r5_w2vec/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3232f7d47d486442c966a5ede893ce491146c91c12183e41ad1502aaa794f94e
+size 321

data/weights/r5_w2vec/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:232950462057e8844d6cd2bba6ac9f1d875309d7d17277aba3645174f0768cd3
+size 4866847

data/weights/r5_w2vec/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:695093769fefcccc7a42b38ef0c1b6a032945aafbd50e4809115cd55d4195ac7
+size 198

data/weights/r7_w2vec/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c1e0224de6018d8ddb966f853fc8c08e5d8eef5b7057b7542d82a9ae4befca40
+size 298

data/weights/r7_w2vec/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:14fbfaeb91e6066f59c894638fe7ba5ca9d019298d5ba80e29d7584dab0afee6
+size 4866847

data/weights/r7_w2vec/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f731ec381f0dea7d7fe99936aad0dceb4a442af251093c78116c2889ae1db64
+size 22245

data/weights/r9_local/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f7b79f12cbbf692de41095bff5cfeab906e97efe18343c74125e91748812ca3
+size 320

data/weights/r9_local/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1bf0fa150c7cb85123497d5cf784386f7351665b196c635b26856ae044d52151
+size 35323

data/weights/r9_local/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:973371a1053fd9a84b93188111cb732ae6610abaf3dae65aaac31b27e96db7b1
+size 21912

data/weights/r9_nbrs/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66acb2a59e43e975f09dd0acfaebf214ef822dfb4c2e04df71946988f0a1ba63
+size 320

data/weights/r9_nbrs/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a515d0cdc564b3db7c2357ea8a0b1495fd482dd83d4ebd407616c4a65f9d8cbf
+size 35323

data/weights/r9_nbrs/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2fb3010bc18246eb2eb81e9e670583a507d4497601e2939bf94825e868a0484c
+size 21912

data/weights/r9_w2vec/args.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:373b05bda0a2c96a04500f850a143a6bc164c16744f6c1efad3bbfe0fc2a3261
+size 298

data/weights/r9_w2vec/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3a53428d08e1f68cffdb78f6499d8fae7d42f7878f65d4e56ce6f0eeb37fa3a
+size 4866847

data/weights/r9_w2vec/train.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf2141e8f42e65470cd391b93f3dfa9ce8d82b6f30cb27eb4b8cbb60d3744e83
+size 22279

models.py

@@ -3,6 +3,7 @@ from torch import nn
 import torch.nn.functional as F
 from torch import Tensor
 from typing import Optional, List
+from timm.models.layers import trunc_normal_
 
 
 class LayerNorm(nn.Module):
@@ -303,6 +304,11 @@ class UNetEncoder(nn.Module):
         x = self.final(x)
 
         return x
+    
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=.02)
+            nn.init.constant_(m.bias, 0)
 
 
 class Decoder(nn.Module):

requirements.txt

@@ -6,3 +6,4 @@ torch
 numpy
 shiny
 shinyswatch
+timm

utils.py

@@ -1,5 +1,11 @@
+import argparse
 import numpy as np
 import pickle
+import os
+import yaml
+import torch
+import torch.nn as nn
+from models import UNetEncoder, Decoder
 
 
 def load_training_data(
@@ -50,3 +56,55 @@ def load_training_data(
         return C, names, Y, M
     else:
         return C, names, Y, M, data["pp_locs"]
+
+
+def radius_from_dir(s: str, prefix: str):
+    return int(s.split("/")[-1].split("_")[0].replace(prefix, ""))
+
+
+def load_models(dirs: dict, prefix="h", nd=5):
+    D = {}
+    for name, datadir in dirs.items():
+        radius = radius_from_dir(datadir, prefix)
+        args = argparse.Namespace()
+        with open(os.path.join(datadir, "args.yaml"), "r") as io:
+            for k, v in yaml.load(io, Loader=yaml.FullLoader).items():
+                setattr(args, k, v)
+                if k == "nbrs_av":
+                    setattr(args, "av_nbrs", v)
+                elif k == "av_nbrs":
+                    setattr(args, "nbrs_av", v)
+
+        bn_type ="frn" if not hasattr(args, "bn_type") else args.bn_type
+        mkw = dict(
+            n_hidden=args.nhidden,
+            depth=args.depth,
+            num_res=args.nres,
+            ksize=args.ksize,
+            groups=args.groups,
+            batchnorm=True,
+            batchnorm_type=bn_type,
+        )
+
+        dkw = dict(batchnorm=True, offset=True, batchnorm_type=bn_type)
+        dev = "cuda" if torch.cuda.is_available() else "cpu"
+        if not args.local and args.nbrs_av == 0:
+            enc = UNetEncoder(nd, args.nhidden, **mkw).to(dev)
+            dec = Decoder(args.nhidden, nd, args.nhidden, **dkw).to(dev)
+        else:
+            enc = nn.Identity()
+            dec = Decoder(nd, nd, args.nhidden, **dkw).to(dev)
+        mod = nn.ModuleDict({"enc": enc, "dec": dec})
+        objs = dict(
+            mod=mod,
+            args=args,
+            radius=radius,
+            nbrs_av=args.nbrs_av,
+            local=args.local,
+        )
+        mod.eval()
+        for p in mod.parameters():
+            p.requires_grad = False
+        mod = mod.to(dev)
+        D[datadir] = objs
+    return D