First model version

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Elijah Cole
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,12 +1,73 @@
----
-title: Sinr
-emoji: 🏃
-colorFrom: green
-colorTo: red
-sdk: gradio
-sdk_version: 3.38.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Spatial Implicit Neural Representations for Global-Scale Species Mapping - ICML 2023
+
+Code for training and evaluating global-scale species range estimation models. This code enables the recreation of the results from our ICML 2023 paper [Spatial Implicit Neural Representations for Global-Scale Species Mapping](https://arxiv.org/abs/2306.02564). 
+
+## 🌍 Overview 
+Estimating the geographical range of a species from sparse observations is a challenging and important geospatial prediction problem. Given a set of locations where a species has been observed, the goal is to build a model to predict whether the species is present or absent at any location. In this work, we use Spatial Implicit Neural Representations (SINRs) to jointly estimate the geographical range of thousands of species simultaneously. SINRs scale gracefully, making better predictions as we increase the number of training species and the amount of training data per species. We introduce four new range estimation and spatial representation learning benchmarks, and we use them to demonstrate that noisy and biased crowdsourced data can be combined with implicit neural representations to approximate expert-developed range maps for many species.
+
+![Model Prediction](images/sinr_traverse.gif)
+<sup>Above we visualize predictions from one of our SINR models trained on data from [iNaturalist](inaturalist.org). On the left we show the learned species embedding space, where each point represents a different species. On the right we see the predicted range of the species corresponding to the red dot on the left.<sup>
+
+## 🔍 Getting Started 
+
+#### Installing Required Packages
+
+1. We recommend using an isolated Python environment to avoid dependency issues. Install the Anaconda Python 3.9 distribution for your operating system from [here](https://www.anaconda.com/download). 
+
+2. Create a new environment and activate it:
+```bash
+ conda create -y --name sinr_icml python==3.9
+ conda activate sinr_icml
+```
+
+3. After activating the environment, install the required packages:
+```bash
+ pip3 install -r requirements.txt
+```
+
+#### Data Download and Preparation
+Instructions for downloading the data in `data/README.md`.
+
+## 🗺️ Generating Predictions
+To generate predictions for a model in the form of an image, run the following command: 
+```bash
+ python viz_map.py --taxa_id 130714
+```
+Here, `--taxa_id` is the id number for a species of interest from [iNaturalist](https://www.inaturalist.org/taxa/130714). If you want to generate predictions for a random species, add the `--rand_taxa` instead. 
+
+Note, before you run this command you need to first download the data as described in `data/README.md`. In addition, if you want to evaluate some of the pretrained models from the paper, you need to download those first and place them at `sinr/pretrained_models`. See `web_app/README.md` for more details.
+
+There is also an interactive browser-based demo available in `web_app`.
+
+## 🚅 Training and Evaluating Models
+
+To train and evaluate a model, run the following command:
+```bash
+ python train_and_evaluate_models.py
+```
+
+#### Hyperparameters
+Common parameters of interest can be set within `train_and_evaluate_models.py`. All other parameters are exposed in `setup.py`. 
+
+#### Outputs
+By default, trained models and evaluation results will be saved to a folder in the `experiments` directory. Evaluation results will also be printed to the command line. 
+
+#### Interactive Model Visualizer
+To visualize range predictions from pretrained SINR models, please follow the instructions in `web_app/README.md`. 
+
+##  🙏 Acknowledgements
+This project was enabled by data from the Cornell Lab of Ornithology, The International Union for the Conservation of Nature, iNaturalist, NASA, USGS, JAXA, CIESIN, and UC Merced. We are especially indebted to the [iNaturalist](inaturalist.org) and [eBird](https://ebird.org) communities for their data collection efforts. We also thank Matt Stimas-Mackey and Sam Heinrich for their help with data curation. This project was funded by the [Climate Change AI Innovation Grants](https://www.climatechange.ai/blog/2022-04-13-innovation-grants) program, hosted by Climate Change AI with the support of the Quadrature Climate Foundation, Schmidt Futures, and the Canada Hub of Future Earth. This work was also supported by the Caltech Resnick Sustainability Institute and an NSF Graduate Research Fellowship (grant number DGE1745301).  
+
+If you find our work useful in your research please consider citing our paper.  
+```
+@inproceedings{SINR_icml23,
+  title     = {{Spatial Implicit Neural Representations for Global-Scale Species Mapping}},
+  author    = {Cole, Elijah and Van Horn, Grant and Lange, Christian and Shepard, Alexander and Leary, Patrick and Perona, Pietro and Loarie, Scott and Mac Aodha, Oisin},
+  booktitle = {ICML},
+  year = {2023}
+}
+```
+
+## 📜 Disclaimer
+Extreme care should be taken before making any decisions based on the outputs of models presented here. Our goal in this work is to demonstrate the promise of large-scale representation learning for species range estimation, not to provide definitive range maps. Our models are trained on biased data and have not been calibrated or validated beyond the experiments illustrated in the paper. 
+

app.py

@@ -0,0 +1,180 @@
+import gradio as gr
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import json
+import os
+import torch
+
+import utils
+import models
+import datasets
+
+
+def load_taxa_metadata(file_path):
+    taxa_names_file = open(file_path, "r")
+    data = taxa_names_file.read().split("\n")
+    data = [dd for dd in data if dd != '']
+    taxa_ids = []
+    taxa_names = []
+    for tt in range(len(data)):
+        id, nm = data[tt].split('\t')
+        taxa_ids.append(int(id))
+        taxa_names.append(nm)
+    taxa_names_file.close()
+    return dict(zip(taxa_ids, taxa_names))
+ 
+    
+def generate_prediction(taxa_id, selected_model, settings, threshold):
+            
+    # select the model to use
+    if selected_model == 'AN_FULL max 10':
+        model_path = 'pretrained_models/model_an_full_input_enc_sin_cos_hard_cap_num_per_class_10.pt'
+    elif selected_model == 'AN_FULL max 100':
+        model_path = 'pretrained_models/model_an_full_input_enc_sin_cos_hard_cap_num_per_class_100.pt'
+    elif selected_model == 'AN_FULL max 1000':
+        model_path = 'pretrained_models/model_an_full_input_enc_sin_cos_hard_cap_num_per_class_1000.pt'
+    elif selected_model == 'Distilled env model':
+        model_path = 'pretrained_models/model_an_full_input_enc_sin_cos_distilled_from_env.pt'
+
+    # load params
+    with open('paths.json', 'r') as f:
+        paths = json.load(f)
+                
+    # configs
+    eval_params = {}
+    eval_params['device'] = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    eval_params['model_path'] = model_path
+    eval_params['taxa_id'] = int(taxa_id)
+    eval_params['rand_taxa'] = 'Random taxa' in settings
+    eval_params['set_max_cmap_to_1'] = False   
+    eval_params['disable_ocean_mask'] = 'Distilled env model' in settings
+    eval_params['threshold'] = threshold if 'Threshold' in settings else -1.0
+        
+    # load model
+    train_params = torch.load(eval_params['model_path'], map_location='cpu')
+    model = models.get_model(train_params['params'])
+    model.load_state_dict(train_params['state_dict'], strict=True)
+    model = model.to(eval_params['device'])
+    model.eval()
+    if train_params['params']['input_enc'] in ['env', 'sin_cos_env']:
+        raster = datasets.load_env(norm=train_params['params']['env_norm'])
+    else:
+        raster = None
+    enc = utils.CoordEncoder(train_params['params']['input_enc'], raster=raster)
+
+    # user specified random taxa
+    if eval_params['rand_taxa']: 
+        print('Selecting random taxa')
+        eval_params['taxa_id'] = np.random.choice(train_params['params']['class_to_taxa'])
+
+    # load taxa of interest 
+    if eval_params['taxa_id'] in train_params['params']['class_to_taxa']:
+        class_of_interest = train_params['params']['class_to_taxa'].index(eval_params['taxa_id'])
+    else:
+        print(f'Error: Taxa specified that is not in the model: {eval_params["taxa_id"]}')
+        fig = plt.figure()
+        plt.imshow(np.zeros((1,1)), vmin=0, vmax=1.0, cmap=plt.cm.plasma)
+        plt.axis('off')
+        plt.tight_layout()    
+        op_html = f'<h2><a href="https://www.inaturalist.org/taxa/{eval_params["taxa_id"]}" target="_blank">{eval_params["taxa_id"]}</a></h2> Error: specified taxa is not in the model.'    
+            
+        return op_html, fig, eval_params['taxa_id']
+    print(f'Loading taxa: {eval_params["taxa_id"]}')
+
+    # load ocean mask
+    mask = np.load(os.path.join(paths['masks'], 'ocean_mask.npy'))
+    mask_inds = np.where(mask.reshape(-1) == 1)[0]
+
+    # generate input features
+    locs = utils.coord_grid(mask.shape)
+    if not eval_params['disable_ocean_mask']:
+        locs = locs[mask_inds, :]
+    locs = torch.from_numpy(locs)
+    locs_enc = enc.encode(locs).to(eval_params['device'])
+
+    # make prediction
+    with torch.no_grad():
+        preds = model(locs_enc, return_feats=False, class_of_interest=class_of_interest).cpu().numpy()
+
+    # threshold predictions
+    if eval_params['threshold'] > 0:
+        print(f'Applying threshold of {eval_params["threshold"]} to the predictions.')
+        preds[preds<eval_params['threshold']] = 0.0
+        preds[preds>=eval_params['threshold']] = 1.0
+
+    # mask data
+    if not eval_params['disable_ocean_mask']:
+        op_im = np.ones((mask.shape[0] * mask.shape[1])) * np.nan  # set to NaN
+        op_im[mask_inds] = preds
+    else:
+        op_im = preds
+
+    # reshape and create masked array for visualization
+    op_im = op_im.reshape((mask.shape[0], mask.shape[1]))
+    op_im = np.ma.masked_invalid(op_im) 
+    if eval_params['set_max_cmap_to_1']:
+        vmax = 1.0
+    else:
+        vmax = np.max(op_im)
+
+    # set color for masked values
+    cmap = plt.cm.plasma
+    cmap.set_bad(color='none')  
+
+    plt.rcParams['figure.figsize'] = 24,12
+    fig = plt.figure()
+    plt.imshow(op_im, vmin=0, vmax=vmax, cmap=cmap)
+    plt.axis('off')
+    plt.tight_layout()
+    
+    # generate html for ouput display
+    taxa_name_str = taxa_names[eval_params['taxa_id']]
+    op_html = f'<h2><a href="https://www.inaturalist.org/taxa/{eval_params["taxa_id"]}" target="_blank">{taxa_name_str}</a></h2> (click for more info)'    
+    return op_html, fig, gr.Number.update(value=eval_params['taxa_id'])
+
+
+# load metadata
+taxa_names = load_taxa_metadata('taxa_02_08_2023_names.txt')  
+
+
+with gr.Blocks(title="SINR Demo") as demo:
+    top_text = "Visualization code to explore species range predictions "\
+               "from Spatial Implicit Neural Representation (SINR) models from "\
+               "[our](https://arxiv.org/abs/2306.02564) ICML 2023 paper." 
+    gr.Markdown("# SINR Visualization Demo")
+    gr.Markdown(top_text)
+    
+    with gr.Row():
+        selected_taxa = gr.Number(label="Taxa ID", value=130714)
+        select_model = gr.Dropdown(["AN_FULL max 10", "AN_FULL max 100", "AN_FULL max 1000", "Distilled env model"],
+                                    value="AN_FULL max 1000", label="Model")
+    with gr.Row():
+        settings = gr.CheckboxGroup(["Random taxa", "Disable ocean mask", "Threshold"], label="Settings")
+        threshold = gr.Slider(0, 1, 0, label="Threshold")
+
+    with gr.Row():
+        submit_button = gr.Button("Run Model")
+
+    with gr.Row():
+        output_text = gr.HTML(label="Species Name:")
+        
+    with gr.Row():
+        output_image = gr.Plot(label="Predicted occupancy")
+
+    end_text = "**Note:** Extreme care should be taken before making any decisions "\
+               "based on the outputs of models presented here. "\
+               "The goal of this work is to demonstrate the promise of large-scale "\
+               "representation learning for species range estimation.  "\
+               "Our models are trained on biased data and have not been calibrated "\
+               "or validated beyondthe experiments illustrated in the paper."
+    gr.Markdown(end_text)
+                
+    submit_button.click(
+        fn = generate_prediction,
+        inputs=[selected_taxa, select_model, settings, threshold],
+        outputs=[output_text, output_image, selected_taxa]
+    )
+
+demo.launch()  

data/masks/ocean_mask.npy

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

datasets.py

@@ -0,0 +1,194 @@
+import os
+import numpy as np
+import json
+import pandas as pd
+from calendar import monthrange
+import torch
+import utils
+
+class LocationDataset(torch.utils.data.Dataset):
+    def __init__(self, locs, labels, classes, class_to_taxa, input_enc, device):
+
+        # handle input encoding:
+        self.input_enc = input_enc
+        if self.input_enc in ['env', 'sin_cos_env']:
+            raster = load_env()
+        else:
+            raster = None
+        self.enc = utils.CoordEncoder(input_enc, raster)
+        
+        # define some properties:
+        self.locs = locs
+        self.loc_feats = self.enc.encode(self.locs) 
+        self.labels = labels
+        self.classes = classes
+        self.class_to_taxa = class_to_taxa
+        
+        # useful numbers:
+        self.num_classes = len(np.unique(labels))
+        self.input_dim = self.loc_feats.shape[1]
+
+        if self.enc.raster is not None:
+            self.enc.raster = self.enc.raster.to(device)
+
+    def __len__(self):
+        return self.loc_feats.shape[0]
+
+    def __getitem__(self, index):
+        loc_feat  = self.loc_feats[index, :]
+        loc       = self.locs[index, :]
+        class_id  = self.labels[index]
+        return loc_feat, loc, class_id
+
+def load_env():
+    with open('paths.json', 'r') as f:
+        paths = json.load(f)
+    raster = load_context_feats(os.path.join(paths['env'],'bioclim_elevation_scaled.npy'))
+    return raster
+
+def load_context_feats(data_path):
+    context_feats = np.load(data_path).astype(np.float32)
+    context_feats = torch.from_numpy(context_feats)
+    return context_feats
+
+def load_inat_data(ip_file, taxa_of_interest=None):
+
+    print('\nLoading  ' + ip_file)
+    data = pd.read_csv(ip_file)
+
+    # remove outliers
+    num_obs = data.shape[0]
+    data = data[((data['latitude'] <= 90) & (data['latitude'] >= -90) & (data['longitude'] <= 180) & (data['longitude'] >= -180) )]
+    if (num_obs - data.shape[0]) > 0:
+        print(num_obs - data.shape[0], 'items filtered due to invalid locations')
+
+    if 'accuracy' in data.columns:
+        data.drop(['accuracy'], axis=1, inplace=True)
+
+    if 'positional_accuracy' in data.columns:
+        data.drop(['positional_accuracy'], axis=1, inplace=True)
+        
+    if 'geoprivacy' in data.columns:
+        data.drop(['geoprivacy'], axis=1, inplace=True)
+
+    if 'observed_on' in data.columns:
+        data.rename(columns = {'observed_on':'date'}, inplace=True)
+
+    num_obs_orig = data.shape[0]
+    data = data.dropna()
+    size_diff = num_obs_orig - data.shape[0]
+    if size_diff > 0:
+        print(size_diff, 'observation(s) with a NaN entry out of' , num_obs_orig, 'removed')
+    
+    # keep only taxa of interest:
+    if taxa_of_interest is not None:
+        num_obs_orig = data.shape[0]
+        data = data[data['taxon_id'].isin(taxa_of_interest)]
+        print(num_obs_orig - data.shape[0], 'observation(s) out of' , num_obs_orig, 'from different taxa removed')
+
+    print('Number of unique classes {}'.format(np.unique(data['taxon_id'].values).shape[0]))
+
+    locs = np.vstack((data['longitude'].values, data['latitude'].values)).T.astype(np.float32)
+    taxa = data['taxon_id'].values.astype(np.int)
+
+    if 'user_id' in data.columns:
+        users = data['user_id'].values.astype(np.int)
+        _, users = np.unique(users, return_inverse=True)
+    elif 'observer_id' in data.columns:
+        users = data['observer_id'].values.astype(np.int)
+        _, users = np.unique(users, return_inverse=True)
+    else:
+        users = np.ones(taxa.shape[0], dtype=np.int)*-1
+
+    # Note - assumes that dates are in format YYYY-MM-DD
+    years  = np.array([int(d_str[:4])   for d_str in data['date'].values])
+    months = np.array([int(d_str[5:7])  for d_str in data['date'].values])
+    days   = np.array([int(d_str[8:10]) for d_str in data['date'].values])
+    days_per_month = np.cumsum([0] + [monthrange(2018, mm)[1] for mm in range(1, 12)])
+    dates  = days_per_month[months-1] + days-1
+    dates  = np.round((dates) / 364.0, 4).astype(np.float32)
+    if 'id' in data.columns:
+        obs_ids = data['id'].values
+    elif 'observation_uuid' in data.columns:
+        obs_ids = data['observation_uuid'].values
+    
+    return locs, taxa, users, dates, years, obs_ids
+
+def choose_aux_species(current_species, num_aux_species, aux_species_seed):
+    if num_aux_species == 0:
+        return []
+    with open('paths.json', 'r') as f:
+        paths = json.load(f)
+    data_dir = paths['train']
+    taxa_file = os.path.join(data_dir, 'geo_prior_train_meta.json')
+    with open(taxa_file, 'r') as f:
+        inat_large_metadata = json.load(f)
+    aux_species_candidates = [x['taxon_id'] for x in inat_large_metadata]
+    aux_species_candidates = np.setdiff1d(aux_species_candidates, current_species)
+    print(f'choosing {num_aux_species} species to add from {len(aux_species_candidates)} candidates')
+    rng = np.random.default_rng(aux_species_seed)
+    idx_rand_aux_species = rng.permutation(len(aux_species_candidates))
+    aux_species = list(aux_species_candidates[idx_rand_aux_species[:num_aux_species]])
+    return aux_species
+
+def get_taxa_of_interest(species_set='all', num_aux_species=0, aux_species_seed=123, taxa_file_snt=None):
+    if species_set == 'all':
+        return None
+    if species_set == 'snt_birds':
+        assert taxa_file_snt is not None
+        with open(taxa_file_snt, 'r') as f: #
+            taxa_subsets = json.load(f)
+        taxa_of_interest = list(taxa_subsets['snt_birds'])
+    else:
+        raise NotImplementedError
+    # optionally add some other species back in: 
+    aux_species = choose_aux_species(taxa_of_interest, num_aux_species, aux_species_seed)
+    taxa_of_interest.extend(aux_species)
+    return taxa_of_interest
+
+def get_idx_subsample_observations(labels, hard_cap=-1, hard_cap_seed=123):
+    if hard_cap == -1:
+        return np.arange(len(labels))
+    print(f'subsampling (up to) {hard_cap} per class for the training set')
+    class_counts = {id: 0 for id in np.unique(labels)}
+    ss_rng = np.random.default_rng(hard_cap_seed)
+    idx_rand = ss_rng.permutation(len(labels))
+    idx_ss = []
+    for i in idx_rand:
+        class_id = labels[i]
+        if class_counts[class_id] < hard_cap:
+            idx_ss.append(i)
+            class_counts[class_id] += 1
+    idx_ss = np.sort(idx_ss)
+    print(f'final training set size: {len(idx_ss)}')
+    return idx_ss
+
+def get_train_data(params):
+    with open('paths.json', 'r') as f:
+        paths = json.load(f)
+    data_dir = paths['train']
+    obs_file  = os.path.join(data_dir, 'geo_prior_train.csv')
+    taxa_file = os.path.join(data_dir, 'geo_prior_train_meta.json')
+    taxa_file_snt = os.path.join(data_dir, 'taxa_subsets.json')
+
+    taxa_of_interest = get_taxa_of_interest(params['species_set'], params['num_aux_species'], params['aux_species_seed'], taxa_file_snt)
+
+    locs, labels, _, _, _, _ = load_inat_data(obs_file, taxa_of_interest)
+    unique_taxa, class_ids = np.unique(labels, return_inverse=True)
+    class_to_taxa = unique_taxa.tolist()
+
+    # load class names
+    class_info_file = json.load(open(taxa_file, 'r'))
+    class_names_file = [cc['latin_name'] for cc in class_info_file]
+    taxa_ids_file = [cc['taxon_id'] for cc in class_info_file]
+    classes = dict(zip(taxa_ids_file, class_names_file))
+    
+    idx_ss = get_idx_subsample_observations(labels, params['hard_cap_num_per_class'], params['hard_cap_seed'])
+
+    locs = torch.from_numpy(np.array(locs)[idx_ss]) # convert to Tensor
+
+    labels = torch.from_numpy(np.array(class_ids)[idx_ss])
+
+    ds = LocationDataset(locs, labels, classes, class_to_taxa, params['input_enc'], params['device'])
+
+    return ds

eval.py

@@ -0,0 +1,362 @@
+import numpy as np
+import pandas as pd
+import random
+import torch
+import time
+import os
+import copy
+import json
+import tifffile
+import h3
+import setup
+
+from sklearn.linear_model import RidgeCV
+from sklearn.preprocessing import MinMaxScaler
+from sklearn.metrics import average_precision_score
+
+import utils
+import models
+import datasets
+
+class EvaluatorSNT:
+    def __init__(self, train_params, eval_params):
+        self.train_params = train_params
+        self.eval_params = eval_params
+        with open('paths.json', 'r') as f:
+            paths = json.load(f)
+        D = np.load(os.path.join(paths['snt'], 'snt_res_5.npy'), allow_pickle=True)
+        D = D.item()
+        self.loc_indices_per_species = D['loc_indices_per_species']
+        self.labels_per_species = D['labels_per_species']
+        self.taxa = D['taxa']
+        self.obs_locs = D['obs_locs']
+        self.obs_locs_idx = D['obs_locs_idx']
+    
+    def get_labels(self, species):
+        species = str(species)
+        lat = []
+        lon = []
+        gt = []
+        for hx in self.data:
+            cur_lat, cur_lon = h3.h3_to_geo(hx)
+            if species in self.data[hx]:
+                cur_label = int(len(self.data[hx][species]) > 0)
+                gt.append(cur_label)
+                lat.append(cur_lat)
+                lon.append(cur_lon)
+        lat = np.array(lat).astype(np.float32)
+        lon = np.array(lon).astype(np.float32)
+        obs_locs = np.vstack((lon, lat)).T
+        gt = np.array(gt).astype(np.float32)
+        return obs_locs, gt
+
+    def run_evaluation(self, model, enc):
+        results = {}
+
+        # set seeds:
+        np.random.seed(self.eval_params['seed'])
+        random.seed(self.eval_params['seed'])
+
+        # evaluate the geo model for each taxon
+        results['mean_average_precision'] = np.zeros((len(self.taxa)), dtype=np.float32)
+        # get eval locations and apply input encoding
+        obs_locs = torch.from_numpy(self.obs_locs).to(self.eval_params['device'])
+        loc_feat = enc.encode(obs_locs)
+        # get classes to eval
+        classes_of_interest = np.array([np.where(np.array(self.train_params['class_to_taxa']) == tt)[0] for tt in self.taxa]).squeeze()
+        classes_of_interest = torch.from_numpy(classes_of_interest)
+        # generate model predictions for classes of interest at eval locations
+        with torch.no_grad():
+            loc_emb = model(loc_feat, return_feats=True)
+            wt = model.class_emb.weight[classes_of_interest, :]
+            pred_mtx = torch.matmul(loc_emb, wt.T).cpu().numpy()
+
+        split_rng = np.random.default_rng(self.eval_params['split_seed'])
+
+        for tt_id, tt in enumerate(self.taxa):
+            # generate ground truth labels for current taxa
+            cur_class_of_interest = np.where(self.taxa == tt)[0][0]
+            cur_loc_indices = np.array(self.loc_indices_per_species[cur_class_of_interest])
+            cur_labels = np.array(self.labels_per_species[cur_class_of_interest])
+        
+            # apply per-species split:
+            assert self.eval_params['split'] in ['all', 'val', 'test']
+            if self.eval_params['split'] != 'all':
+                num_val = np.floor(len(cur_labels) * self.eval_params['val_frac']).astype(int)
+                idx_rand = split_rng.permutation(len(cur_labels))
+                if self.eval_params['split'] == 'val':
+                    idx_sel = idx_rand[:num_val]
+                elif self.eval_params['split'] == 'test':
+                    idx_sel = idx_rand[num_val:]
+                cur_loc_indices = cur_loc_indices[idx_sel]
+                cur_labels = cur_labels[idx_sel]
+
+            # extract model predictions for current taxa from prediction matrix:
+            pred = pred_mtx[cur_loc_indices, tt_id]
+
+            # compute the AP for each taxa
+            results['mean_average_precision'][tt_id] = average_precision_score((cur_labels > 0).astype(np.int32), pred)
+
+
+        valid_taxa = ~np.isnan(results['mean_average_precision'])
+
+        # store results
+        results['per_species_average_precision_all'] = copy.deepcopy(results['mean_average_precision'])
+        per_species_average_precision_valid = results['per_species_average_precision_all'][valid_taxa]
+        results['mean_average_precision'] = per_species_average_precision_valid.mean()
+        results['num_eval_species_w_valid_ap'] = valid_taxa.sum()
+        results['num_eval_species_total'] = len(self.taxa)
+
+        return results
+
+    def report(self, results):
+        for field in ['mean_average_precision', 'num_eval_species_w_valid_ap', 'num_eval_species_total']:
+            print(f'{field}: {results[field]}')
+
+class EvaluatorIUCN:
+
+    def __init__(self, train_params, eval_params):
+        self.train_params = train_params
+        self.eval_params = eval_params
+        with open('paths.json', 'r') as f:
+            paths = json.load(f)
+        with open(os.path.join(paths['iucn'], 'iucn_res_5.json'), 'r') as f:
+            self.data = json.load(f)
+        self.obs_locs = np.array(self.data['locs'], dtype=np.float32)
+        self.taxa = [int(tt) for tt in self.data['taxa_presence'].keys()]
+    
+    def run_evaluation(self, model, enc):
+        results = {}
+
+        results['per_species_average_precision_all'] = np.zeros(len(self.taxa), dtype=np.float32)
+        # get eval locations and apply input encoding
+        obs_locs = torch.from_numpy(self.obs_locs).to(self.eval_params['device'])
+        loc_feat = enc.encode(obs_locs)
+        
+        # get classes to eval
+        classes_of_interest = torch.from_numpy(np.array([np.where(np.array(self.train_params['class_to_taxa']) == tt)[0] for tt in self.taxa]).squeeze())
+        with torch.no_grad():
+            # generate model predictions for classes of interest at eval locations
+            loc_emb = model(loc_feat, return_feats=True)
+            wt = model.class_emb.weight[classes_of_interest, :]
+            pred_mtx = torch.matmul(loc_emb, wt.T)
+
+        for tt_id, tt in enumerate(self.taxa):
+            class_of_interest = np.where(np.array(self.train_params['class_to_taxa']) == tt)[0]
+
+            if len(class_of_interest) == 0:
+                # taxa of interest is not in the model
+                pred = None
+            else:
+                # extract model predictions for current taxa from prediction matrix
+                pred = pred_mtx[:, tt_id]
+            
+            # evaluate accuracy
+            if pred is None:
+                results['per_species_average_precision_all'][tt_id] = np.nan
+            else:
+                gt = np.zeros(obs_locs.shape[0], dtype=np.float32)
+                gt[self.data['taxa_presence'][str(tt)]] = 1.0
+                # average precision score:
+                results['per_species_average_precision_all'][tt_id] = average_precision_score(gt, pred)
+        
+        valid_taxa = ~np.isnan(results['per_species_average_precision_all'])
+
+        # store results
+        per_species_average_precision_valid = results['per_species_average_precision_all'][valid_taxa]
+        results['mean_average_precision'] = per_species_average_precision_valid.mean()
+        results['num_eval_species_w_valid_ap'] = valid_taxa.sum()
+        results['num_eval_species_total'] = len(self.taxa)
+        return results
+
+    def report(self, results):
+        for field in ['mean_average_precision', 'num_eval_species_w_valid_ap', 'num_eval_species_total']:
+            print(f'{field}: {results[field]}')
+
+class EvaluatorGeoPrior:
+
+    def __init__(self, train_params, eval_params):
+        # store parameters:
+        self.train_params = train_params
+        self.eval_params = eval_params
+        with open('paths.json', 'r') as f:
+            paths = json.load(f)
+        # load vision model predictions:
+        self.data = np.load(os.path.join(paths['geo_prior'], 'geo_prior_model_preds.npz'))
+        print('\n', self.data['probs'].shape[0], 'total test observations')
+        # load locations:
+        meta = pd.read_csv(os.path.join(paths['geo_prior'], 'geo_prior_model_meta.csv'))
+        self.obs_locs  = np.vstack((meta['longitude'].values, meta['latitude'].values)).T.astype(np.float32)
+        # taxonomic mapping:
+        self.taxon_map = self.find_mapping_between_models(self.data['model_to_taxa'], self.train_params['class_to_taxa'])
+        print(self.taxon_map.shape[0], 'out of', len(self.data['model_to_taxa']), 'taxa in both vision and geo models')
+
+    def find_mapping_between_models(self, vision_taxa, geo_taxa):
+        # this will output an array of size N_overlap X 2
+        # the first column will be the indices of the vision model, and the second is their
+        # corresponding index in the geo model
+        taxon_map = np.ones((vision_taxa.shape[0], 2), dtype=np.int32)*-1
+        taxon_map[:, 0] = np.arange(vision_taxa.shape[0])
+        geo_taxa_arr = np.array(geo_taxa)
+        for tt_id, tt in enumerate(vision_taxa):
+            ind = np.where(geo_taxa_arr==tt)[0]
+            if len(ind) > 0:
+                taxon_map[tt_id, 1] = ind[0]
+        inds = np.where(taxon_map[:, 1]>-1)[0]
+        taxon_map = taxon_map[inds, :]
+        return taxon_map
+    
+    def convert_to_inat_vision_order(self, geo_pred_ip, vision_top_k_prob, vision_top_k_inds, vision_taxa, taxon_map):
+        # this is slow as we turn the sparse input back into the same size as the dense one
+        vision_pred = np.zeros((geo_pred_ip.shape[0], len(vision_taxa)), dtype=np.float32)
+        geo_pred = np.ones((geo_pred_ip.shape[0], len(vision_taxa)), dtype=np.float32)
+        vision_pred[np.arange(vision_pred.shape[0])[..., np.newaxis], vision_top_k_inds] = vision_top_k_prob
+
+        geo_pred[:, taxon_map[:, 0]] = geo_pred_ip[:, taxon_map[:, 1]]
+
+        return geo_pred, vision_pred
+    
+    def run_evaluation(self, model, enc):
+        results = {}
+
+        # loop over in batches
+        batch_start = np.hstack((np.arange(0, self.data['probs'].shape[0], self.eval_params['batch_size']), self.data['probs'].shape[0]))
+        correct_pred = np.zeros(self.data['probs'].shape[0])
+
+        print('\nbid\t w geo\t wo geo')
+        for bb_id, bb in enumerate(range(len(batch_start)-1)):
+            batch_inds = np.arange(batch_start[bb], batch_start[bb+1])
+
+            vision_probs = self.data['probs'][batch_inds, :]
+            vision_inds = self.data['inds'][batch_inds, :]
+            gt = self.data['labels'][batch_inds]
+
+            obs_locs_batch = torch.from_numpy(self.obs_locs[batch_inds, :]).to(self.eval_params['device'])
+            loc_feat = enc.encode(obs_locs_batch)
+
+            with torch.no_grad():
+                geo_pred = model(loc_feat).cpu().numpy()
+
+            geo_pred, vision_pred = self.convert_to_inat_vision_order(geo_pred, vision_probs, vision_inds,
+                                                                 self.data['model_to_taxa'], self.taxon_map)
+
+            comb_pred = np.argmax(vision_pred*geo_pred, 1)
+            comb_pred = (comb_pred==gt)
+            correct_pred[batch_inds] = comb_pred
+        
+        results['vision_only_top_1'] = float((self.data['inds'][:, -1] == self.data['labels']).mean())
+        results['vision_geo_top_1'] = float(correct_pred.mean())
+        return results
+    
+    def report(self, results):
+        print('\nOverall accuracy vision only model', round(results['vision_only_top_1'], 3))
+        print('Overall accuracy of geo model     ', round(results['vision_geo_top_1'], 3))
+        print('Gain                              ', round(results['vision_geo_top_1'] - results['vision_only_top_1'], 3))
+
+class EvaluatorGeoFeature:
+
+    def __init__(self, train_params, eval_params):
+        self.train_params = train_params
+        self.eval_params = eval_params
+        with open('paths.json', 'r') as f:
+            paths = json.load(f)
+        self.data_path = paths['geo_feature']
+        self.country_mask = tifffile.imread(os.path.join(paths['masks'], 'USA_MASK.tif')) == 1
+        self.raster_names = ['ABOVE_GROUND_CARBON', 'ELEVATION', 'LEAF_AREA_INDEX', 'NON_TREE_VEGITATED', 'NOT_VEGITATED', 'POPULATION_DENSITY', 'SNOW_COVER', 'SOIL_MOISTURE', 'TREE_COVER']
+        self.raster_names_log_transform = ['POPULATION_DENSITY']
+
+    def load_raster(self, raster_name, log_transform=False):
+        raster = tifffile.imread(os.path.join(self.data_path, raster_name + '.tif')).astype(np.float32)
+        valid_mask = ~np.isnan(raster).copy() & self.country_mask
+        # log scaling:
+        if log_transform:
+            raster[valid_mask] = np.log1p(raster[valid_mask] - raster[valid_mask].min())
+        # 0/1 scaling:
+        raster[valid_mask] -= raster[valid_mask].min()
+        raster[valid_mask] /= raster[valid_mask].max()
+        
+        return raster, valid_mask
+    
+    def get_split_labels(self, raster, split_ids, split_of_interest):
+        # get the GT labels for a subset
+        inds_y, inds_x = np.where(split_ids==split_of_interest) 
+        return raster[inds_y, inds_x]
+    
+    def get_split_feats(self, model, enc, split_ids, split_of_interest):
+        locs = utils.coord_grid(self.country_mask.shape, split_ids=split_ids, split_of_interest=split_of_interest)
+        locs = torch.from_numpy(locs).to(self.eval_params['device'])
+        locs_enc = enc.encode(locs)
+        with torch.no_grad():
+            feats = model(locs_enc, return_feats=True).cpu().numpy()
+        return feats
+    
+    def run_evaluation(self, model, enc):
+        results = {}
+        for raster_name in self.raster_names:
+            do_log_transform = raster_name in self.raster_names_log_transform
+            raster, valid_mask = self.load_raster(raster_name, do_log_transform)
+            split_ids = utils.create_spatial_split(raster, valid_mask, cell_size=self.eval_params['cell_size'])
+            feats_train = self.get_split_feats(model, enc, split_ids=split_ids, split_of_interest=1)
+            feats_test = self.get_split_feats(model, enc, split_ids=split_ids, split_of_interest=2)
+            labels_train = self.get_split_labels(raster, split_ids, 1)
+            labels_test = self.get_split_labels(raster, split_ids, 2)
+            scaler = MinMaxScaler()
+            feats_train_scaled = scaler.fit_transform(feats_train)
+            feats_test_scaled = scaler.transform(feats_test)
+            clf = RidgeCV(alphas=(0.1, 1.0, 10.0), normalize=False, cv=10, fit_intercept=True, scoring='r2').fit(feats_train_scaled, labels_train)
+            train_score = clf.score(feats_train_scaled, labels_train)
+            test_score = clf.score(feats_test_scaled, labels_test)
+            results[f'train_r2_{raster_name}'] = float(train_score)
+            results[f'test_r2_{raster_name}'] = float(test_score)
+            results[f'alpha_{raster_name}'] = float(clf.alpha_)
+        return results
+    
+    def report(self, results):
+        report_fields = [x for x in results if 'test_r2' in x]
+        for field in report_fields:
+            print(f'{field}: {results[field]}')
+        print(np.mean([results[field] for field in report_fields]))
+
+def launch_eval_run(overrides):
+
+    eval_params = setup.get_default_params_eval(overrides)
+
+    # set up model:
+    eval_params['model_path'] = os.path.join(eval_params['exp_base'], eval_params['experiment_name'], eval_params['ckp_name'])
+    train_params = torch.load(eval_params['model_path'], map_location='cpu')
+    model = models.get_model(train_params['params'])
+    model.load_state_dict(train_params['state_dict'], strict=True)
+    model = model.to(eval_params['device'])
+    model.eval()
+
+    # create input encoder:
+    if train_params['params']['input_enc'] in ['env', 'sin_cos_env']:
+        raster = datasets.load_env().to(eval_params['device'])
+    else:
+        raster = None
+    enc = utils.CoordEncoder(train_params['params']['input_enc'], raster=raster)
+
+    t = time.time()
+    if eval_params['eval_type'] == 'snt':
+        eval_params['split'] = 'test' # val, test, all
+        eval_params['val_frac'] = 0.50
+        eval_params['split_seed'] = 7499
+        evaluator = EvaluatorSNT(train_params['params'], eval_params)
+        results = evaluator.run_evaluation(model, enc)
+        evaluator.report(results)
+    elif eval_params['eval_type'] == 'iucn':
+        evaluator = EvaluatorIUCN(train_params['params'], eval_params)
+        results = evaluator.run_evaluation(model, enc)
+        evaluator.report(results)
+    elif eval_params['eval_type'] == 'geo_prior':
+        evaluator = EvaluatorGeoPrior(train_params['params'], eval_params)
+        results = evaluator.run_evaluation(model, enc)
+        evaluator.report(results)
+    elif eval_params['eval_type'] == 'geo_feature':
+        evaluator = EvaluatorGeoFeature(train_params['params'], eval_params)
+        results = evaluator.run_evaluation(model, enc)
+        evaluator.report(results)
+    else:
+        raise NotImplementedError('Eval type not implemented.')
+    print(f'evaluation completed in {np.around((time.time()-t)/60, 1)} min')
+    return results

losses.py

@@ -0,0 +1,146 @@
+import torch
+import utils
+
+def get_loss_function(params):
+    if params['loss'] == 'an_full':
+        return an_full
+    elif params['loss'] == 'an_slds':
+        return an_slds
+    elif params['loss'] == 'an_ssdl':
+        return an_ssdl
+    elif params['loss'] == 'an_full_me':
+        return an_full_me
+    elif params['loss'] == 'an_slds_me':
+        return an_slds_me
+    elif params['loss'] == 'an_ssdl_me':
+        return an_ssdl_me
+
+def neg_log(x):
+    return -torch.log(x + 1e-5)
+
+def bernoulli_entropy(p):
+    entropy = p * neg_log(p) + (1-p) * neg_log(1-p)
+    return entropy
+
+def an_ssdl(batch, model, params, loc_to_feats, neg_type='hard'):
+
+    inds = torch.arange(params['batch_size'])
+
+    loc_feat, _, class_id = batch
+    loc_feat = loc_feat.to(params['device'])
+    class_id = class_id.to(params['device'])
+    
+    assert model.inc_bias == False
+    batch_size = loc_feat.shape[0]
+    
+    # create random background samples and extract features
+    rand_loc = utils.rand_samples(batch_size, params['device'], rand_type='spherical')
+    rand_feat = loc_to_feats(rand_loc, normalize=False)
+    
+    # get location embeddings
+    loc_cat = torch.cat((loc_feat, rand_feat), 0)  # stack vertically
+    loc_emb_cat = model(loc_cat, return_feats=True)
+    loc_emb = loc_emb_cat[:batch_size, :]
+    loc_emb_rand = loc_emb_cat[batch_size:, :]
+    
+    loc_pred = torch.sigmoid(model.class_emb(loc_emb))
+    loc_pred_rand = torch.sigmoid(model.class_emb(loc_emb_rand))
+    
+    # data loss
+    loss_pos = neg_log(loc_pred[inds[:batch_size], class_id])
+    if neg_type == 'hard':
+        loss_bg = neg_log(1.0 - loc_pred_rand[inds[:batch_size], class_id]) # assume negative
+    elif neg_type == 'entropy':
+        loss_bg = -1 * bernoulli_entropy(1.0 - loc_pred_rand[inds[:batch_size], class_id]) # entropy
+    else:
+        raise NotImplementedError
+    
+    # total loss
+    loss = loss_pos.mean() + loss_bg.mean()
+    
+    return loss
+
+def an_slds(batch, model, params, loc_to_feats, neg_type='hard'):
+
+    inds = torch.arange(params['batch_size'])
+
+    loc_feat, _, class_id = batch
+    loc_feat = loc_feat.to(params['device'])
+    class_id = class_id.to(params['device'])
+    
+    assert model.inc_bias == False
+    batch_size = loc_feat.shape[0]
+
+    loc_emb = model(loc_feat, return_feats=True)
+    
+    loc_pred = torch.sigmoid(model.class_emb(loc_emb))
+    
+    num_classes = loc_pred.shape[1]
+    bg_class = torch.randint(low=0, high=num_classes-1, size=(batch_size,), device=params['device'])
+    bg_class[bg_class >= class_id[:batch_size]] += 1
+    
+    # data loss
+    loss_pos = neg_log(loc_pred[inds[:batch_size], class_id])
+    if neg_type == 'hard':
+        loss_bg = neg_log(1.0 - loc_pred[inds[:batch_size], bg_class]) # assume negative
+    elif neg_type == 'entropy':
+        loss_bg = -1 * bernoulli_entropy(1.0 - loc_pred[inds[:batch_size], bg_class]) # entropy
+    else:
+        raise NotImplementedError
+    
+    # total loss
+    loss = loss_pos.mean() + loss_bg.mean()
+    
+    return loss
+
+def an_full(batch, model, params, loc_to_feats, neg_type='hard'):
+    
+    inds = torch.arange(params['batch_size'])
+
+    loc_feat, _, class_id = batch
+    loc_feat = loc_feat.to(params['device'])
+    class_id = class_id.to(params['device'])
+    
+    assert model.inc_bias == False
+    batch_size = loc_feat.shape[0]
+    
+    # create random background samples and extract features
+    rand_loc = utils.rand_samples(batch_size, params['device'], rand_type='spherical')
+    rand_feat = loc_to_feats(rand_loc, normalize=False)
+    
+    # get location embeddings
+    loc_cat = torch.cat((loc_feat, rand_feat), 0)  # stack vertically
+    loc_emb_cat = model(loc_cat, return_feats=True)
+    loc_emb = loc_emb_cat[:batch_size, :]
+    loc_emb_rand = loc_emb_cat[batch_size:, :]
+    # get predictions for locations and background locations
+    loc_pred = torch.sigmoid(model.class_emb(loc_emb))
+    loc_pred_rand = torch.sigmoid(model.class_emb(loc_emb_rand))
+    
+    # data loss
+    if neg_type == 'hard':
+        loss_pos = neg_log(1.0 - loc_pred) # assume negative
+        loss_bg = neg_log(1.0 - loc_pred_rand) # assume negative
+    elif neg_type == 'entropy':
+        loss_pos = -1 * bernoulli_entropy(1.0 - loc_pred) # entropy
+        loss_bg = -1 * bernoulli_entropy(1.0 - loc_pred_rand) # entropy
+    else:
+        raise NotImplementedError
+    loss_pos[inds[:batch_size], class_id] = params['pos_weight'] * neg_log(loc_pred[inds[:batch_size], class_id])
+    
+    # total loss
+    loss = loss_pos.mean() + loss_bg.mean()
+    
+    return loss
+
+def an_full_me(batch, model, params, loc_to_feats):
+
+    return an_full(batch, model, params, loc_to_feats, neg_type='entropy')
+
+def an_ssdl_me(batch, model, params, loc_to_feats):
+    
+    return an_ssdl(batch, model, params, loc_to_feats, neg_type='entropy')
+
+def an_slds_me(batch, model, params, loc_to_feats):
+    
+    return an_slds(batch, model, params, loc_to_feats, neg_type='entropy')

models.py

@@ -0,0 +1,85 @@
+import torch
+import torch.utils.data
+import torch.nn as nn
+
+def get_model(params):
+
+    if params['model'] == 'ResidualFCNet':
+        return ResidualFCNet(params['input_dim'], params['num_classes'], params['num_filts'], params['depth'])
+    elif params['model'] == 'LinNet':
+        return LinNet(params['input_dim'], params['num_classes'])
+    else:
+        raise NotImplementedError('Invalid model specified.')
+
+class ResLayer(nn.Module):
+    def __init__(self, linear_size):
+        super(ResLayer, self).__init__()
+        self.l_size = linear_size
+        self.nonlin1 = nn.ReLU(inplace=True)
+        self.nonlin2 = nn.ReLU(inplace=True)
+        self.dropout1 = nn.Dropout()
+        self.w1 = nn.Linear(self.l_size, self.l_size)
+        self.w2 = nn.Linear(self.l_size, self.l_size)
+
+    def forward(self, x):
+        y = self.w1(x)
+        y = self.nonlin1(y)
+        y = self.dropout1(y)
+        y = self.w2(y)
+        y = self.nonlin2(y)
+        out = x + y
+        return out
+
+class ResidualFCNet(nn.Module):
+
+    def __init__(self, num_inputs, num_classes, num_filts, depth=4):
+        super(ResidualFCNet, self).__init__()
+        self.inc_bias = False
+        self.class_emb = nn.Linear(num_filts, num_classes, bias=self.inc_bias)
+        layers = []
+        layers.append(nn.Linear(num_inputs, num_filts))
+        layers.append(nn.ReLU(inplace=True))
+        for i in range(depth):
+            layers.append(ResLayer(num_filts))
+        self.feats = torch.nn.Sequential(*layers)
+
+    def forward(self, x, class_of_interest=None, return_feats=False):
+        loc_emb = self.feats(x)
+        if return_feats:
+            return loc_emb
+        if class_of_interest is None:
+            class_pred = self.class_emb(loc_emb)
+        else:
+            class_pred = self.eval_single_class(loc_emb, class_of_interest)
+        return torch.sigmoid(class_pred)
+
+    def eval_single_class(self, x, class_of_interest):
+        if self.inc_bias:
+            return torch.matmul(x, self.class_emb.weight[class_of_interest, :].T) + self.class_emb.bias[class_of_interest]
+        else:
+            return torch.matmul(x, self.class_emb.weight[class_of_interest, :].T)
+
+class LinNet(nn.Module):
+    def __init__(self, num_inputs, num_classes):
+        super(LinNet, self).__init__()
+        self.num_layers = 0
+        self.inc_bias = False
+        self.class_emb = nn.Linear(num_inputs, num_classes, bias=self.inc_bias)
+        self.feats = nn.Identity()  # does not do anything
+
+    def forward(self, x, class_of_interest=None, return_feats=False):
+        loc_emb = self.feats(x)
+        if return_feats:
+            return loc_emb
+        if class_of_interest is None:
+            class_pred = self.class_emb(loc_emb)
+        else:
+            class_pred = self.eval_single_class(loc_emb, class_of_interest)
+
+        return torch.sigmoid(class_pred)
+
+    def eval_single_class(self, x, class_of_interest):
+        if self.inc_bias:
+            return torch.matmul(x, self.class_emb.weight[class_of_interest, :].T) + self.class_emb.bias[class_of_interest]
+        else:
+            return torch.matmul(x, self.class_emb.weight[class_of_interest, :].T)

paths.json

@@ -0,0 +1,9 @@
+{
+    "masks": "data/masks/",
+    "env": "data/env/", 
+    "train": "data/train/", 
+    "geo_prior": "data/eval/geo_prior/",
+    "snt": "data/eval/snt/", 
+    "iucn": "data/eval/iucn/", 
+    "geo_feature": "data/eval/geo_feature/"
+} 

pretrained_models/model_an_full_input_enc_sin_cos_distilled_from_env.pt

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

pretrained_models/model_an_full_input_enc_sin_cos_hard_cap_num_per_class_10.pt

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

pretrained_models/model_an_full_input_enc_sin_cos_hard_cap_num_per_class_100.pt

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

pretrained_models/model_an_full_input_enc_sin_cos_hard_cap_num_per_class_1000.pt

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

requirements.txt

@@ -0,0 +1,6 @@
+gradio==3.36.1
+h3==3.7.6
+matplotlib==3.7.1
+numpy==1.25.0
+pandas==2.0.3
+torch==1.12.1

setup.py

@@ -0,0 +1,91 @@
+import copy
+import torch
+
+def apply_overrides(params, overrides):
+    params = copy.deepcopy(params)
+    for param_name in overrides:
+        if param_name not in params:
+            print(f'override failed: no parameter named {param_name}')
+            raise ValueError
+        params[param_name] = overrides[param_name]
+    return params
+
+def get_default_params_train(overrides={}):
+
+    params = {}
+    
+    '''
+    misc
+    '''
+    params['device'] = 'cuda' # cuda, cpu
+    params['save_base'] = './experiments/'
+    params['experiment_name'] = 'demo'
+    params['timestamp'] = False
+    
+    '''
+    data
+    '''
+    params['species_set'] = 'all' # all, snt_birds
+    params['hard_cap_seed'] = 9472
+    params['hard_cap_num_per_class'] = -1 # -1 for no hard capping
+    params['aux_species_seed'] = 8099
+    params['num_aux_species'] = 0 # for snt_birds case, how many other species to add in
+
+    '''
+    model
+    '''
+    params['model'] = 'ResidualFCNet'  # ResidualFCNet, LinNet
+    params['num_filts'] = 256  # embedding dimension
+    params['input_enc'] = 'sin_cos' # sin_cos, env, sin_cos_env
+    params['depth'] = 4
+    
+    '''
+    loss
+    '''
+    params['loss'] = 'an_full' # an_full, an_ssdl, an_slds
+    params['pos_weight'] = 2048
+    
+    '''
+    optimization
+    '''
+    params['batch_size'] = 2048
+    params['lr'] = 0.0005
+    params['lr_decay'] = 0.98
+    params['num_epochs'] = 10
+    
+    '''
+    saving
+    '''
+    params['log_frequency'] = 512
+    
+    params = apply_overrides(params, overrides)
+    
+    return params
+
+def get_default_params_eval(overrides={}):
+
+    params = {}
+
+    '''
+    misc
+    '''
+    params['device'] = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    params['seed'] = 2022
+    params['exp_base'] = './experiments'
+    params['ckp_name'] = 'model.pt'
+    params['eval_type'] = 'snt' # snt, iucn, geo_prior, geo_feature
+    params['experiment_name'] = 'demo'
+
+    '''
+    geo prior
+    '''
+    params['batch_size'] = 2048
+
+    '''
+    geo feature
+    '''
+    params['cell_size'] = 25
+
+    params = apply_overrides(params, overrides)
+
+    return params

utils.py

@@ -0,0 +1,143 @@
+import torch
+import numpy as np
+import math 
+import datetime
+
+class CoordEncoder:
+    
+    def __init__(self, input_enc, raster=None):
+        self.input_enc = input_enc    
+        self.raster = raster
+
+    def encode(self, locs, normalize=True):
+        # assumes lon, lat in range [-180, 180] and [-90, 90]
+        if normalize:
+            locs = normalize_coords(locs)
+        if self.input_enc == 'sin_cos': # sinusoidal encoding
+            loc_feats = encode_loc(locs) 
+        elif self.input_enc == 'env': # bioclim variables
+            loc_feats = bilinear_interpolate(locs, self.raster) 
+        elif self.input_enc == 'sin_cos_env': # sinusoidal encoding & bioclim variables
+            loc_feats = encode_loc(locs) 
+            context_feats = bilinear_interpolate(locs, self.raster)
+            loc_feats = torch.cat((loc_feats, context_feats), 1)
+        else:
+            raise NotImplementedError('Unknown input encoding.')
+        return loc_feats
+
+def normalize_coords(locs):
+    # locs is in lon {-180, 180}, lat {90, -90}
+    # output is in the range [-1, 1]
+
+    locs[:,0] /= 180.0
+    locs[:,1] /= 90.0
+
+    return locs
+
+def encode_loc(loc_ip, concat_dim=1):
+    # assumes inputs location are in range -1 to 1
+    # location is lon, lat
+    feats = torch.cat((torch.sin(math.pi*loc_ip), torch.cos(math.pi*loc_ip)), concat_dim)
+    return feats
+
+def bilinear_interpolate(loc_ip, data, remove_nans_raster=True):
+    # loc is N x 2 vector, where each row is [lon,lat] entry
+    #   each entry spans range [-1,1]
+    # data is H x W x C, height x width x channel data matrix
+    # op will be N x C matrix of interpolated features
+
+    assert data is not None
+
+    # map to [0,1], then scale to data size
+    loc = (loc_ip.clone() + 1) / 2.0
+    loc[:,1] = 1 - loc[:,1] # this is because latitude goes from +90 on top to bottom while
+                            # longitude goes from -90 to 90 left to right
+
+    assert not torch.any(torch.isnan(loc))
+    
+    if remove_nans_raster:
+        data[torch.isnan(data)] = 0.0 # replace with mean value (0 is mean post-normalization)
+
+    # cast locations into pixel space
+    loc[:, 0] *= (data.shape[1]-1)
+    loc[:, 1] *= (data.shape[0]-1)
+
+    loc_int = torch.floor(loc).long()  # integer pixel coordinates
+    xx = loc_int[:, 0]
+    yy = loc_int[:, 1]
+    xx_plus = xx + 1
+    xx_plus[xx_plus > (data.shape[1]-1)] = data.shape[1]-1
+    yy_plus = yy + 1
+    yy_plus[yy_plus > (data.shape[0]-1)] = data.shape[0]-1
+
+    loc_delta = loc - torch.floor(loc)   # delta values
+    dx = loc_delta[:, 0].unsqueeze(1)
+    dy = loc_delta[:, 1].unsqueeze(1)
+
+    interp_val = data[yy, xx, :]*(1-dx)*(1-dy) + data[yy, xx_plus, :]*dx*(1-dy) + \
+                 data[yy_plus, xx, :]*(1-dx)*dy   + data[yy_plus, xx_plus, :]*dx*dy
+
+    return interp_val
+
+def rand_samples(batch_size, device, rand_type='uniform'):
+    # randomly sample background locations
+
+    if rand_type == 'spherical':
+        rand_loc = torch.rand(batch_size, 2).to(device)
+        theta1 = 2.0*math.pi*rand_loc[:, 0]
+        theta2 = torch.acos(2.0*rand_loc[:, 1] - 1.0)
+        lat = 1.0 - 2.0*theta2/math.pi
+        lon = (theta1/math.pi) - 1.0
+        rand_loc = torch.cat((lon.unsqueeze(1), lat.unsqueeze(1)), 1)
+
+    elif rand_type == 'uniform':
+        rand_loc = torch.rand(batch_size, 2).to(device)*2.0 - 1.0
+
+    return rand_loc
+
+def get_time_stamp():
+    cur_time = str(datetime.datetime.now())
+    date, time = cur_time.split(' ')
+    h, m, s = time.split(':')
+    s = s.split('.')[0]
+    time_stamp = '{}-{}-{}-{}'.format(date, h, m, s)
+    return time_stamp
+
+def coord_grid(grid_size, split_ids=None, split_of_interest=None):
+    # generate a grid of locations spaced evenly in coordinate space
+
+    feats = np.zeros((grid_size[0], grid_size[1], 2), dtype=np.float32)
+    mg = np.meshgrid(np.linspace(-180, 180, feats.shape[1]), np.linspace(90, -90, feats.shape[0]))
+    feats[:, :, 0] = mg[0]
+    feats[:, :, 1] = mg[1]
+    if split_ids is None or split_of_interest is None:
+        # return feats for all locations
+        # this will be an N x 2 array
+        return feats.reshape(feats.shape[0]*feats.shape[1], 2)
+    else:
+        # only select a subset of locations
+        ind_y, ind_x = np.where(split_ids==split_of_interest)
+
+        # these will be N_subset x 2 in size
+        return feats[ind_y, ind_x, :]
+    
+def create_spatial_split(raster, mask, train_amt=1.0, cell_size=25):
+    # generates a checkerboard style train test split
+    # 0 is invalid, 1 is train, and 2 is test 
+    # c_size is units of pixels
+    split_ids = np.ones((raster.shape[0], raster.shape[1]))
+    start = cell_size
+    for ii in np.arange(0, split_ids.shape[0], cell_size): 
+        if start == 0:
+            start = cell_size
+        else:
+            start = 0
+        for jj in np.arange(start, split_ids.shape[1], cell_size*2): 
+            split_ids[ii:ii+cell_size, jj:jj+cell_size] = 2 
+    split_ids = split_ids*mask
+    if train_amt < 1.0:
+        # take a subset of the data
+        tr_y, tr_x = np.where(split_ids==1) 
+        inds = np.random.choice(len(tr_y), int(len(tr_y)*(1.0-train_amt)), replace=False)
+        split_ids[tr_y[inds], tr_x[inds]] = 0
+    return split_ids