Merge branch 'master' of github.com:GilesStrong/pytorch_inferno

Author: GilesStrong

README.md

@@ -53,10 +53,9 @@ If you have used this implementation of INFERNO in your analysis work and wish t
 ```
 @misc{giles_chatham_strong_2021_4597140,  
   author       = {Giles Chatham Strong},  
-  title        = {LUMIN},  
-  month        = mar,  
-  year         = 2021,  
-  note         = {{Please check https://github.com/GilesStrong/pytorch_inferno/graphs/contributors for the full list of contributors}},  
+  title        = {PyTorch INFERNO},  
+  month        = Mar,  
+  year         = 2021,   
   doi          = {10.5281/zenodo.4597140},  
   url          = {https://doi.org/10.5281/zenodo.4597140}  
 }

docs/Gemfile

@@ -3,6 +3,6 @@ source "https://rubygems.org"
 gem 'github-pages', group: :jekyll_plugins
 
 # Added at 2019-11-25 10:11:40 -0800 by jhoward:
-gem "nokogiri", "< 1.11.1"
+gem "nokogiri", "< 1.11.5"
 gem "jekyll", ">= 3.7"
 gem "kramdown", ">= 2.3.0"

docs/Gemfile.lock

@@ -201,14 +201,14 @@ GEM
       rb-fsevent (~> 0.10, >= 0.10.3)
       rb-inotify (~> 0.9, >= 0.9.10)
     mercenary (0.3.6)
-    mini_portile2 (2.5.0)
+    mini_portile2 (2.5.1)
     minima (2.5.1)
       jekyll (>= 3.5, < 5.0)
       jekyll-feed (~> 0.9)
       jekyll-seo-tag (~> 2.1)
     minitest (5.14.1)
     multipart-post (2.1.1)
-    nokogiri (1.11.0)
+    nokogiri (1.11.4)
       mini_portile2 (~> 2.5.0)
       racc (~> 1.4)
     octokit (4.18.0)
@@ -221,7 +221,7 @@ GEM
     rb-fsevent (0.10.4)
     rb-inotify (0.10.1)
       ffi (~> 1.0)
-    rexml (3.2.4)
+    rexml (3.2.5)
     rouge (3.19.0)
     ruby-enum (0.8.0)
       i18n
@@ -257,7 +257,7 @@ DEPENDENCIES
   github-pages
   jekyll (>= 3.7)
   kramdown (>= 2.3.0)
-  nokogiri (< 1.11.1)
+  nokogiri (< 1.11.5)
 
 BUNDLED WITH
    2.0.2

nbs/07_inferno_exact.ipynb

@@ -138,7 +138,7 @@
     "            if hasattr(c, 'loss_is_meaned'): c.loss_is_meaned = False  # Ensure that average losses are correct\n",
     "    \n",
     "    @abstractmethod\n",
-    "    def _get_up_down(self, x_s:Tensor, x_b:Tensor) -> Tuple[Tuple[Optional[Tensor],Optional[Tensor]],Tuple[Optional[Tensor],Optional[Tensor]]]:\n",
+    "    def _get_up_down(self, x_s:Tensor, x_b:Tensor, w_s:Optional[Tensor]=None, w_b:Optional[Tensor]=None) -> Tuple[Tuple[Optional[Tensor],Optional[Tensor]],Tuple[Optional[Tensor],Optional[Tensor]]]:\n",
     "        r'''Compute upd/down shapes for signal and background seperately. Overide this for specific problem.'''\n",
     "        pass\n",
     "    \n",
@@ -163,8 +163,10 @@
     "    def on_forwards_end(self) -> None:\n",
     "        r'''Compute loss and replace wrapper loss value'''\n",
     "        b = self.wrapper.y.squeeze() == 0\n",
-    "        f_s = self.to_shape(self.wrapper.y_pred[~b])\n",
-    "        f_b = self.to_shape(self.wrapper.y_pred[b])\n",
+    "        w_s = self.wrapper.w[~b] if self.wrapper.w is not None else None\n",
+    "        w_b = self.wrapper.w[b] if self.wrapper.w is not None else None\n",
+    "        f_s = self.to_shape(self.wrapper.y_pred[~b], w_s)\n",
+    "        f_b = self.to_shape(self.wrapper.y_pred[b], w_b)\n",
     "        (f_s_up,f_s_dw),(f_b_up,f_b_dw)= self._get_up_down(self.wrapper.x[~b], self.wrapper.x[b])\n",
     "        self.wrapper.loss_val = self.get_ikk(f_s_nom=f_s, f_b_nom=f_b, f_s_up=f_s_up, f_s_dw=f_s_dw, f_b_up=f_b_up, f_b_dw=f_b_dw)"
    ]
@@ -195,7 +197,7 @@
     "            self.l_mod_t[0][0,2] = self.l_mods[0]/self.l_init\n",
     "            self.l_mod_t[1][0,2] = self.l_mods[1]/self.l_init\n",
     "    \n",
-    "    def _get_up_down(self, x_s:Tensor, x_b:Tensor) -> Tuple[Tuple[Optional[Tensor],Optional[Tensor]],Tuple[Optional[Tensor],Optional[Tensor]]]:\n",
+    "    def _get_up_down(self, x_s:Tensor, x_b:Tensor, **kwargs) -> Tuple[Tuple[Optional[Tensor],Optional[Tensor]],Tuple[Optional[Tensor],Optional[Tensor]]]:\n",
     "        if self.r_mods is None and self.l_mods is None: return (None,None),(None,None)\n",
     "        u,d = [],[]\n",
     "        if self.r_mods is not None:\n",

nbs/08_inferno_interp.ipynb

@@ -84,18 +84,23 @@ def get_inv_ikk(self, f_s:Tensor, f_b:Tensor, f_s_asimov:Tensor, f_b_asimov:Tens
         return torch.inverse(h)[self.poi_idx,self.poi_idx]
 
     @staticmethod
-    def to_shape(p:Tensor) -> Tensor:
-        f = p.sum(0)+1e-7
+    def to_shape(p:Tensor, w:Optional[Tensor]=None) -> Tensor:
+        f = (p*w).sum(0)+1e-7 if w is not None else p.sum(0)+1e-7
         return f/f.sum()
 
     def on_forwards_end(self) -> None:
         r'''Compute loss and replace wrapper loss value'''
+
+        w_s = self.wrapper.w[~self.b_mask] if self.wrapper.w is not None else None
+        w_b = self.wrapper.w[self.b_mask] if self.wrapper.w is not None else None
+
         # Shapes with derivatives w.r.t. nuisances
-        f_s = self.to_shape(self.wrapper.y_pred[~self.b_mask])
-        f_b = self.to_shape(self.wrapper.y_pred[self.b_mask])
+        f_s = self.to_shape(self.wrapper.y_pred[~self.b_mask], w_s)
+        f_b = self.to_shape(self.wrapper.y_pred[self.b_mask], w_b)
+
         # Shapes without derivatives w.r.t. nuisances
-        f_s_asimov = self.to_shape(self.wrapper.model(self.wrapper.x[~self.b_mask].detach())) if self.s_shape_alpha else f_s
-        f_b_asimov = self.to_shape(self.wrapper.model(self.wrapper.x[self.b_mask].detach()))  if self.b_shape_alpha else f_b
+        f_s_asimov = self.to_shape(self.wrapper.model(self.wrapper.x[~self.b_mask].detach()), w_s) if self.s_shape_alpha else f_s
+        f_b_asimov = self.to_shape(self.wrapper.model(self.wrapper.x[self.b_mask].detach()), w_b)  if self.b_shape_alpha else f_b
 
         self.wrapper.loss_val = self.get_inv_ikk(f_s=f_s, f_b=f_b, f_s_asimov=f_s_asimov, f_b_asimov=f_b_asimov)
 
@@ -141,7 +146,7 @@ def on_train_begin(self) -> None:
             if hasattr(c, 'loss_is_meaned'): c.loss_is_meaned = False  # Ensure that average losses are correct
 
     @abstractmethod
-    def _get_up_down(self, x_s:Tensor, x_b:Tensor) -> Tuple[Tuple[Optional[Tensor],Optional[Tensor]],Tuple[Optional[Tensor],Optional[Tensor]]]:
+    def _get_up_down(self, x_s:Tensor, x_b:Tensor, w_s:Optional[Tensor]=None, w_b:Optional[Tensor]=None) -> Tuple[Tuple[Optional[Tensor],Optional[Tensor]],Tuple[Optional[Tensor],Optional[Tensor]]]:
         r'''Compute upd/down shapes for signal and background seperately. Overide this for specific problem.'''
         pass
 
@@ -166,8 +171,10 @@ def get_ikk(self, f_s_nom:Tensor, f_b_nom:Tensor, f_s_up:Optional[Tensor], f_s_d
     def on_forwards_end(self) -> None:
         r'''Compute loss and replace wrapper loss value'''
         b = self.wrapper.y.squeeze() == 0
-        f_s = self.to_shape(self.wrapper.y_pred[~b])
-        f_b = self.to_shape(self.wrapper.y_pred[b])
+        w_s = self.wrapper.w[~b] if self.wrapper.w is not None else None
+        w_b = self.wrapper.w[b] if self.wrapper.w is not None else None
+        f_s = self.to_shape(self.wrapper.y_pred[~b], w_s)
+        f_b = self.to_shape(self.wrapper.y_pred[b], w_b)
         (f_s_up,f_s_dw),(f_b_up,f_b_dw)= self._get_up_down(self.wrapper.x[~b], self.wrapper.x[b])
         self.wrapper.loss_val = self.get_ikk(f_s_nom=f_s, f_b_nom=f_b, f_s_up=f_s_up, f_s_dw=f_s_dw, f_b_up=f_b_up, f_b_dw=f_b_dw)
 
@@ -191,8 +198,8 @@ def on_train_begin(self) -> None:
             self.l_mod_t[0][0,2] = self.l_mods[0]/self.l_init
             self.l_mod_t[1][0,2] = self.l_mods[1]/self.l_init
 
-    def _get_up_down(self, x_s:Tensor, x_b:Tensor) -> Tuple[Tuple[Optional[Tensor],Optional[Tensor]],Tuple[Optional[Tensor],Optional[Tensor]]]:
-        if self.r_mods is None and self.l_mods is None: return None,None
+    def _get_up_down(self, x_s:Tensor, x_b:Tensor, **kwargs) -> Tuple[Tuple[Optional[Tensor],Optional[Tensor]],Tuple[Optional[Tensor],Optional[Tensor]]]:
+        if self.r_mods is None and self.l_mods is None: return (None,None),(None,None)
         u,d = [],[]
         if self.r_mods is not None:
             with torch.no_grad(): x_b = x_b+self.r_mod_t[0]