time truncation

Author: victor-shepardson

notepredictor/notepredictor/distributions.py

@@ -1,4 +1,5 @@
 import math
+import numpy as np
 
 import torch
 from torch import nn
@@ -97,33 +98,63 @@ def forward(self, h, x):
         return r
 
     def cdf(self, h, x):
+        """
+        Args:
+            h: Tensor[...,n_params]
+            x: Tensor[...]
+                `h` should broadcast with `x[...,None]`
+        Returns:
+            cdf: Tensor[...] (shape of `x` broadcasted with `h[...,0]`)
+        """
         log_pi, loc, s = self.get_params(h)  
-        x_ = (x[...,None] - loc) * s
-        cdfs = x_.sigmoid()
+        cdfs = self.cdf_components(loc, s, x)
         cdf = (cdfs * log_pi.softmax(-1)).sum(-1)
         return cdf
 
-    def sample(self, h, shape=None):
+    def cdf_components(self, loc, s, x):
+        x_ = (x[...,None] - loc) * s
+        return x_.sigmoid()
+
+    def sample(self, h, truncate=None, shape=None):
         """
         Args:
-            shape: additional sample shape to be prepended to dims
+            h: Tensor[...,n_params]
+            shape: additional sample shape to be prepended to dims or None
+        Returns:
+            Tensor[*shape,...] (h without last dimension, prepended with `shape`)
         """
         if shape is None:
             unwrap = True
             shape = 1
         else:
             unwrap = False
+
+        if truncate is None:
+            truncate = (-np.inf, np.inf)
+        truncate = torch.tensor(truncate)
+
         log_pi, loc, s = self.get_params(h)
         scale = 1/s
 
-        c = D.Categorical(logits=log_pi).sample((shape,))
+        # cdfs: [...,bound,component]
+        cdfs = self.cdf_components(loc[...,None,:], s[...,None,:], truncate) 
+        # prob. mass of each component witin bounds
+        trunc_probs = cdfs[...,1,:] - cdfs[...,0,:] # [...,component]
+        probs = log_pi.exp() * trunc_probs # reweighted mixture component probs
+
+        c = D.Categorical(probs).sample((shape,))
         # move sample dimension first
         loc = loc.movedim(-1, 0).gather(0, c)
         scale = scale.movedim(-1, 0).gather(0, c)
+        upper = cdfs[...,1,:].movedim(-1, 0).gather(0, c)
+        lower = cdfs[...,0,:].movedim(-1, 0).gather(0, c)
 
         u = torch.rand(shape, *h.shape[:-1])
+        # truncate
+        u = u * (upper-lower) + lower
 
-        x = loc + scale * (u.log() - (1 - u).log())
+        # x = loc + scale * (u.log() - (1 - u).log())
+        x = loc - scale * (1/u - 1).log()
         x = x.clamp(self.lo, self.hi)
         return x[0] if unwrap else x
 

notepredictor/notepredictor/model.py

@@ -200,7 +200,8 @@ def cell_state(self):
 
     def get_samplers(self, 
             pitch_topk=None, index_pitch=None, allow_start=False, allow_end=False, 
-            sweep_time=False):
+            sweep_time=False, trunc_time=None):
+
         def sample_pitch(x):
             if not allow_start:
                 x[...,self.start_token] = -np.inf
@@ -214,15 +215,20 @@ def sample_pitch(x):
                 return D.Categorical(logits=x).sample()
 
         def sample_time(x):
+            # TODO: respect trunc_time when sweep_time is True
             if sweep_time:
+                if trunc_time is not None:
+                    raise NotImplementedError("""
+                    trunc_time with sweep_time needs implementation
+                    """)
                 assert x.shape[0]==1, "batch size should be 1 here"
                 log_pi, loc, s = self.time_dist.get_params(x)
                 idx = log_pi.squeeze().argsort()[:9]
                 loc = loc.squeeze()[idx].sort().values[...,None] # multiple times in batch dim
                 # print(loc.shape)
                 return loc
             else:
-                return self.time_dist.sample(x)
+                return self.time_dist.sample(x, truncate=trunc_time)
 
         return (
             sample_pitch, 
@@ -310,7 +316,7 @@ def predict(self,
             pitch, time, vel, 
             fix_pitch=None, fix_time=None, fix_vel=None, 
             pitch_topk=None, index_pitch=None, allow_start=False, allow_end=False,
-            sweep_time=False):
+            sweep_time=False, trunc_time=None):
         """
         consume the most recent note and return a prediction for the next note.
 
@@ -329,6 +335,7 @@ def predict(self,
             allow_end: if False, zero probaility for sampling the end token
             sweep_time: if True, instead of sampling time, choose a diverse set of
                 times and stack along the batch dimension
+            trunc_time: if not None, truncate the time distribution to (lo, hi)
 
         Returns: dict of
             'pitch': int. predicted MIDI number of next note.
@@ -358,7 +365,8 @@ def predict(self,
             modalities = list(zip(
                 self.projections,
                 self.get_samplers(
-                    pitch_topk, index_pitch, allow_start, allow_end, sweep_time),
+                    pitch_topk, index_pitch, allow_start, allow_end, 
+                    sweep_time, trunc_time),
                 self.embeddings,
                 ))
 
@@ -379,7 +387,10 @@ def predict(self,
             det_idx, cons_idx, uncons_idx = [], [], []
             for i,(item, embed) in enumerate(zip(fix, self.embeddings)):
                 if item is None:
-                    if (i==1 and sweep_time) or (i==0 and pitch_topk):
+                    if (
+                        i==1 and (sweep_time or (trunc_time is not None)) or
+                        i==0 and pitch_topk
+                        ):
                         cons_idx.append(i)
                     else:
                         uncons_idx.append(i)
@@ -398,8 +409,9 @@ def predict(self,
             # TODO: allow constraints; 
             # attempt to sort the strongest constraints first
             # constraints can be:
-            # discrete set, in which case evaluate probs and then sample categorical
-            # range, in which case truncate
+            # discrete set, in which case evaluate probs and then sample categorical;
+            # range, in which case truncate;
+            # temperature?
 
             perm_h_tgt = [h_tgt[i] for i in perm]
             while len(undet_idx):