Merge branch 'master' of https://github.com/monash-emu/AuTuMN

.gitignore

@@ -93,4 +93,5 @@ data/pbi_outputs_for_opti/by_location/fully_immune/output_table_fully_immune_by_
 autumn/projects/tuberculosis/marshall_islands/outputs/all_outputs/*
 
 # ignore Romain's remote output runs
-user/rragonnet/remote_run_outputs
+user/rragonnet/remote_run_outputs
+user/rragonnet/temp/

autumn/models/sm_covid2/stratifications/immunity.py

@@ -6,7 +6,7 @@
 from jax import numpy as jnp
 from summer2.parameters import Function, Data, Time
 
-from autumn.model_features.functional import piecewise_constant
+from summer2.functions.util import piecewise_constant
 
 from autumn.core.inputs.database import get_input_db
 from autumn.models.sm_covid2.constants import IMMUNITY_STRATA, ImmunityStratum, FlowName

autumn/projects/sm_covid2/common_school/included_countries.yml

@@ -1,85 +1,56 @@
 all:
-  AGO: Angola
-  ALB: Albania
   ARG: Argentina
-  ARM: Armenia
   AUS: Australia
   AUT: Austria
-  AZE: Azerbaijan
   BEL: Belgium
   BGD: Bangladesh
   BGR: Bulgaria
   BIH: Bosnia and Herzegovina
   BLR: Belarus
   BOL: Bolivia, Plurinational State of
   BRA: Brazil
-  BWA: Botswana
   CAN: Canada
   CHL: Chile
-  CHN: China
-  CMR: Cameroon
-  COD: Congo, The Democratic Republic of the
   COL: Colombia
   CRI: Costa Rica
   CZE: Czechia
   DEU: Germany
   DNK: Denmark
-  DOM: Dominican Republic
-  DZA: Algeria
   ECU: Ecuador
   EGY: Egypt
   ESP: Spain
-  EST: Estonia
-  ETH: Ethiopia
   FIN: Finland
   FRA: France
   GBR: United Kingdom
   GEO: Georgia
-  GHA: Ghana
   GRC: Greece
   GTM: Guatemala
-  GUY: Guyana
   HND: Honduras
   HRV: Croatia
   HUN: Hungary
   IDN: Indonesia
   IND: India
   IRL: Ireland
-  IRN: Iran, Islamic Republic of
   IRQ: Iraq
   ISR: Israel
   ITA: Italy
-  JAM: Jamaica
   JOR: Jordan
   JPN: Japan
   KAZ: Kazakhstan
   KEN: Kenya
-  KHM: Cambodia
   KOR: Korea, Republic of
-  KWT: Kuwait
   LBN: Lebanon
   LKA: Sri Lanka
   LTU: Lithuania
-  LUX: Luxembourg
   LVA: Latvia
   MAR: Morocco
   MDA: Moldova, Republic of
   MEX: Mexico
   MKD: North Macedonia
   MMR: Myanmar
-  MNE: Montenegro
-  MNG: Mongolia
-  MOZ: Mozambique
-  MUS: Mauritius
-  MWI: Malawi
   MYS: Malaysia
-  NAM: Namibia
-  NGA: Nigeria
   NLD: Netherlands
-  NOR: Norway
   NPL: Nepal
-  NZL: New Zealand
-  OMN: Oman
   PAK: Pakistan
   PAN: Panama
   PER: Peru
@@ -89,130 +60,19 @@ all:
   PRY: Paraguay
   ROU: Romania
   RUS: Russian Federation
-  RWA: Rwanda
   SAU: Saudi Arabia
-  SEN: Senegal
-  SGP: Singapore
-  SLV: El Salvador
   SRB: Serbia
-  SUR: Suriname
   SVK: Slovakia
   SVN: Slovenia
   SWE: Sweden
-  SWZ: Eswatini
-  SYR: Syrian Arab Republic
   THA: Thailand
-  TTO: Trinidad and Tobago
-  TUN: Tunisia
   TUR: Turkey
-  UGA: Uganda
   UKR: Ukraine
   URY: Uruguay
   USA: United States
   VEN: Venezuela, Bolivarian Republic of
   VNM: Viet Nam
   ZAF: South Africa
-  ZMB: Zambia
-  ZWE: Zimbabwe
-google_mobility:
-  AGO: Angola
-  ARG: Argentina
-  AUS: Australia
-  AUT: Austria
-  BEL: Belgium
-  BGD: Bangladesh
-  BGR: Bulgaria
-  BIH: Bosnia and Herzegovina
-  BLR: Belarus
-  BOL: Bolivia, Plurinational State of
-  BRA: Brazil
-  BWA: Botswana
-  CAN: Canada
-  CHL: Chile
-  CMR: Cameroon
-  COL: Colombia
-  CRI: Costa Rica
-  CZE: Czechia
-  DEU: Germany
-  DNK: Denmark
-  DOM: Dominican Republic
-  ECU: Ecuador
-  EGY: Egypt
-  ESP: Spain
-  EST: Estonia
-  FIN: Finland
-  FRA: France
-  GBR: United Kingdom
-  GEO: Georgia
-  GHA: Ghana
-  GRC: Greece
-  GTM: Guatemala
-  HND: Honduras
-  HRV: Croatia
-  HUN: Hungary
-  IDN: Indonesia
-  IND: India
-  IRL: Ireland
-  IRQ: Iraq
-  ISR: Israel
-  ITA: Italy
-  JAM: Jamaica
-  JOR: Jordan
-  JPN: Japan
-  KAZ: Kazakhstan
-  KEN: Kenya
-  KHM: Cambodia
-  KOR: Korea, Republic of
-  KWT: Kuwait
-  LBN: Lebanon
-  LKA: Sri Lanka
-  LTU: Lithuania
-  LUX: Luxembourg
-  LVA: Latvia
-  MAR: Morocco
-  MDA: Moldova, Republic of
-  MEX: Mexico
-  MKD: North Macedonia
-  MMR: Myanmar
-  MNG: Mongolia
-  MOZ: Mozambique
-  MUS: Mauritius
-  MYS: Malaysia
-  NGA: Nigeria
-  NLD: Netherlands
-  NOR: Norway
-  NPL: Nepal
-  NZL: New Zealand
-  OMN: Oman
-  PAK: Pakistan
-  PAN: Panama
-  PER: Peru
-  PHL: Philippines
-  POL: Poland
-  PRT: Portugal
-  PRY: Paraguay
-  ROU: Romania
-  RUS: Russian Federation
-  RWA: Rwanda
-  SAU: Saudi Arabia
-  SEN: Senegal
-  SGP: Singapore
-  SLV: El Salvador
-  SRB: Serbia
-  SVK: Slovakia
-  SVN: Slovenia
-  SWE: Sweden
-  THA: Thailand
-  TTO: Trinidad and Tobago
-  TUR: Turkey
-  UGA: Uganda
-  UKR: Ukraine
-  URY: Uruguay
-  USA: United States
-  VEN: Venezuela, Bolivarian Republic of
-  VNM: Viet Nam
-  ZAF: South Africa
-  ZMB: Zambia
   ZWE: Zimbabwe
 national_sero:
   AUS: Australia
@@ -234,7 +94,6 @@ national_sero:
   HRV: Croatia
   HUN: Hungary
   IND: India
-  IRN: Iran, Islamic Republic of
   ISR: Israel
   ITA: Italy
   JOR: Jordan
@@ -245,17 +104,10 @@ national_sero:
   LBN: Lebanon
   LTU: Lithuania
   MEX: Mexico
-  MWI: Malawi
-  NGA: Nigeria
-  NOR: Norway
   NPL: Nepal
-  OMN: Oman
   PAK: Pakistan
   PRT: Portugal
-  SEN: Senegal
-  SGP: Singapore
   SVN: Slovenia
   SWE: Sweden
   USA: United States
   ZAF: South Africa
-  ZMB: Zambia

autumn/projects/sm_covid2/common_school/output_plots/country_spec.py

@@ -208,7 +208,7 @@ def plot_two_scenarios(axis, uncertainty_dfs, output_name, iso3, include_unc=Fal
     axis.set_ylim((0, plot_ymax))
 
     if include_legend:
-        axis.legend()
+        axis.legend(title="(median and IQR)")
 
     axis.yaxis.set_major_formatter(tick.FuncFormatter(y_fmt))
 
@@ -505,12 +505,12 @@ def test_tiling_plot():
     from pathlib import Path
     import pandas as pd
 
-    directory = Path.cwd() / "autumn" / "projects" / "sm_covid2" / "common_school" / "output_plots" / "test_tiling_plot"
-
-    iso3 = "FRA"
-    uncertainty_df = pd.read_parquet(directory / "uncertainty_df.parquet")
+    iso3 = "ARG"
+    directory = Path.cwd() / "user" / "rragonnet" / "remote_run_outputs" / "31747883_full_analysis_26Sep2023_main" / iso3
+    uncertainty_dfs = {sc: pd.read_parquet(directory / f"uncertainty_df_{sc}.parquet") for sc in ["baseline", "scenario_1"]}
     diff_quantiles_df = pd.read_parquet(directory / "diff_quantiles_df.parquet")
-    output_folder = directory
-    make_country_output_tiling(iso3, uncertainty_df, diff_quantiles_df, output_folder)
+
+    output_folder = Path.cwd() / "user" / "rragonnet" / "temp"
+    make_country_output_tiling(iso3, uncertainty_dfs, diff_quantiles_df, output_folder)
 
 # test_tiling_plot()

autumn/projects/sm_covid2/common_school/output_plots/multicountry.py

@@ -4,18 +4,24 @@
 
 import plotly.graph_objects as go
 
+YLAB_LOOKUP = {
+   "cases_averted_relative": "% infections averted by school closure", 
+   "deaths_averted_relative": "% deaths averted by school closure",
+   "delta_hospital_peak_relative": "Relative reduction in peak hospital occupancy (%)"
+}
+
+YLAB_LOOKUP_SPLIT = {
+   "cases_averted_relative": "% infections averted<br>by school closure", 
+   "deaths_averted_relative": "% deaths averted<br>by school closure",
+   "delta_hospital_peak_relative": "Relative reduction in<br>peak hospital occupancy (%)"
+}
 
 def plot_multic_relative_outputs(output_dfs_dict: dict[str, pd.DataFrame], req_outputs=["cases_averted_relative", "deaths_averted_relative", "delta_hospital_peak_relative"]):
     n_subplots = len(req_outputs)
     fig, axes = plt.subplots(n_subplots, 1, figsize=(25, n_subplots*6))
 
     this_iso3_list = list(output_dfs_dict.keys())
     n_countries = len(this_iso3_list)
-    ylab_lookup = {
-        "cases_averted_relative": "% infections averted by school closure", 
-        "deaths_averted_relative": "% deaths averted by school closure",
-        "delta_hospital_peak_relative": "Relative reduction in peak hospital occupancy (%)"
-    }
 
     box_width = .4
     med_color = 'white'
@@ -56,8 +62,8 @@ def plot_multic_relative_outputs(output_dfs_dict: dict[str, pd.DataFrame], req_o
 
         axis.set_xticks(ticks=range(1, n_countries + 1), labels=sorted_iso3_list, rotation=90, fontsize=13)
 
-        y_label = ylab_lookup[output]
-        axis.set_ylabel(y_label, fontsize=13)
+        y_label = YLAB_LOOKUP[output]
+        axis.set_ylabel(y_label, fontsize=15)
 
         # add coloured backgorund patches
         xmin, xmax = axis.get_xlim()
@@ -80,6 +86,8 @@ def plot_relative_map(output_dfs_dict: dict[str, pd.DataFrame], req_output="delt
    values = [- 100 * output_dfs_dict[iso3][req_output].loc[0.5] for iso3 in this_iso3_list]
    data_df = pd.DataFrame.from_dict({"iso3": this_iso3_list, "values": values})
 
+   legend_title = YLAB_LOOKUP_SPLIT[req_output]
+
    fig = go.Figure(
       data=go.Choropleth(
          locations=data_df["iso3"], 
@@ -89,7 +97,7 @@ def plot_relative_map(output_dfs_dict: dict[str, pd.DataFrame], req_output="delt
                     [1, 'blue']],  # "Plotly3",
          marker_line_color='darkgrey',
          marker_line_width=0.5,
-         colorbar_title="Relative reduction in<br>peak hospital occupancy",
+         colorbar_title=legend_title,
          colorbar_ticksuffix="%",
       )  
    )

autumn/projects/sm_covid2/common_school/utils.py

@@ -2,15 +2,6 @@
 import os
 from autumn.settings.folders import PROJECTS_PATH
 
-
-def get_school_country_list(nationbal_only=False):
-    source = os.path.join(PROJECTS_PATH, "sm_covid2", "common_school", "included_countries.yml")
-    country_dict = yaml.load(open(source), Loader=yaml.UnsafeLoader)
-    primary_key = "national" if nationbal_only else "all"
-
-    return country_dict[primary_key]   
-
-
 SCHOOL_ISO2_LIST = [
     'AW', 'AF', 'AO', 'AL', 'AE', 'AR', 'AM', 'AG', 'AU', 'AT', 'AZ', 'BI', 'BE', 'BJ', 'BF', 'BD', 'BG', 'BH', 'BS', 
     'BA', 'BY', 'BZ', 'BO', 'BR', 'BB', 'BN', 'BW', 'CF', 'CA', 'CH', 'CL', 'CN', 'CI', 'CM', 'CD', 'CG', 'CO', 'KM',