Updated thinning algo according to NHP's Algorithm 2

easy_tpp/model/torch_model/torch_thinning.py

@@ -1,5 +1,6 @@
 import torch
 import torch.nn as nn
+from easy_tpp.utils import logger
 
 
 class EventSampler(nn.Module):
@@ -34,6 +35,11 @@ def __init__(self, num_sample, num_exp, over_sample_rate, num_samples_boundary,
 
  def compute_intensity_upper_bound(self, time_seq, time_delta_seq, event_seq, intensity_fn,
  compute_last_step_only):
+ # logger.critical(f'time_seq: {time_seq}')
+ # logger.critical(f'time_delta_seq: {time_delta_seq}')
+ # logger.critical(f'event_seq: {event_seq}')
+ # logger.critical(f'intensity_fn: {intensity_fn}')
+ # logger.critical(f'compute_last_step_only: {compute_last_step_only}')
  """Compute the upper bound of intensity at each event timestamp.
 
  Args:
@@ -54,10 +60,10 @@ def compute_intensity_upper_bound(self, time_seq, time_delta_seq, event_seq, int
  steps=self.num_samples_boundary,
  device=self.device)[None, None, :]
 
- # [batch_size, seq_len, num_sample]
+ # [batch_size, seq_len, num_samples_boundary]
  dtime_for_bound_sampled = time_delta_seq[:, :, None] * time_for_bound_sampled
 
- # [batch_size, seq_len, num_sample, event_num]
+ # [batch_size, seq_len, num_samples_boundary, event_num]
  intensities_for_bound = intensity_fn(time_seq,
  time_delta_seq,
  event_seq,
@@ -120,34 +126,46 @@ def sample_uniform_distribution(self, intensity_upper_bound):
 
  return unif_numbers
 
- def sample_accept(self, unif_numbers, sample_rate, total_intensities):
+ def sample_accept(self, unif_numbers, sample_rate, total_intensities, exp_numbers):
  """Do the sample-accept process.
 
- For each parallel draw, find its min criterion： if that < 1.0, the 1st (i.e. smallest) sampled time
- with cri < 1.0 is accepted; if none is accepted, use boundary / maxsampletime for that draw
+ For the accumulated exp (delta) samples drawn for each event timestamp, find (from left to right) the first
+ that makes the criterion < 1 and accept it as the sampled next-event time. If all exp samples are rejected 
+ (criterion >= 1), then we set the sampled next-event time dtime_max.
 
  Args:
  unif_numbers (tensor): [batch_size, max_len, num_sample, num_exp], sampled uniform random number.
  sample_rate (tensor): [batch_size, max_len], sample rate (intensity).
  total_intensities (tensor): [batch_size, seq_len, num_sample, num_exp]
+ exp_numbers (tensor): [batch_size, seq_len, num_sample, num_exp]: sampled exp numbers (delta in Algorithm 2).
 
  Returns:
- list: two tensors,
- criterion, [batch_size, max_len, num_sample, num_exp]
- who_has_accepted_times, [batch_size, max_len, num_sample]
+ result (tensor): [batch_size, seq_len, num_sample], sampled next-event times.
  """
 
  # [batch_size, max_len, num_sample, num_exp]
  criterion = unif_numbers * sample_rate[:, :, None, None] / total_intensities
-
+
+ # [batch_size, max_len, num_sample, num_exp]
+ masked_crit_less_than_1 = torch.where(criterion<1,1,0)
+
  # [batch_size, max_len, num_sample]
- min_cri_each_draw, _ = criterion.min(dim=-1)
-
- # find out unif_numbers * sample_rate < intensity
+ non_accepted_filter = (1-masked_crit_less_than_1).all(dim=3)
+
  # [batch_size, max_len, num_sample]
- who_has_accepted_times = min_cri_each_draw < 1.0
-
- return criterion, who_has_accepted_times
+ first_accepted_indexer = masked_crit_less_than_1.argmax(dim=3)
+
+ # [batch_size, max_len, num_sample,1]
+ # indexer must be unsqueezed to 4D to match the number of dimensions of exp_numbers
+ result_non_accepted_unfiltered = torch.gather(exp_numbers, 3, first_accepted_indexer.unsqueeze(3))
+
+ # [batch_size, max_len, num_sample,1]
+ result = torch.where(non_accepted_filter.unsqueeze(3), torch.tensor(self.dtime_max), result_non_accepted_unfiltered)
+
+ # [batch_size, max_len, num_sample]
+ result = result.squeeze(dim=-1)
+
+ return result
 
  def draw_next_time_one_step(self, time_seq, time_delta_seq, event_seq, dtime_boundary,
  intensity_fn, compute_last_step_only=False):
@@ -177,7 +195,8 @@ def draw_next_time_one_step(self, time_seq, time_delta_seq, event_seq, dtime_bou
  # we apply fast approximation, i.e., re-use exp sample times for computation
  # [batch_size, seq_len, num_exp]
  exp_numbers = self.sample_exp_distribution(intensity_upper_bound)
-
+ exp_numbers = torch.cumsum(exp_numbers, dim=-1)
+
  # 3. compute intensity at sampled times from exp distribution
  # [batch_size, seq_len, num_exp, event_num]
  intensities_at_sampled_times = intensity_fn(time_seq,
@@ -193,46 +212,20 @@ def draw_next_time_one_step(self, time_seq, time_delta_seq, event_seq, dtime_bou
  # add one dim of num_sample: re-use the intensity for samples for prediction
  # [batch_size, seq_len, num_sample, num_exp]
  total_intensities = torch.tile(total_intensities[:, :, None, :], [1, 1, self.num_sample, 1])
+
  # [batch_size, seq_len, num_sample, num_exp]
  exp_numbers = torch.tile(exp_numbers[:, :, None, :], [1, 1, self.num_sample, 1])
-
+ 
  # 4. draw uniform distribution
  # [batch_size, seq_len, num_sample, num_exp]
  unif_numbers = self.sample_uniform_distribution(intensity_upper_bound)
 
  # 5. find out accepted intensities
- # criterion, [batch_size, max_len, num_sample, num_exp]
- # who_has_accepted_times, [batch_size, max_len, num_sample]
- criterion, who_has_accepted_times = self.sample_accept(unif_numbers, intensity_upper_bound,
- total_intensities)
-
- # 6. find out accepted dtimes
- sampled_dtimes_accepted = exp_numbers.clone()
-
- # for unaccepted, use boundary/maxsampletime for that draw
- sampled_dtimes_accepted[criterion >= 1.0] = exp_numbers.max() + 1.0
-
- accepted_times_each_draw, accepted_id_each_draw = sampled_dtimes_accepted.min(dim=-1)
-
- # 7. fill out result
- dtime_boundary_ = dtime_boundary[:, -1:] if compute_last_step_only else dtime_boundary
-
- # [batch_size, seq_len, num_sample]
- dtime_boundary_ = torch.tile(dtime_boundary_[..., None], [1, 1, self.num_sample])
-
  # [batch_size, seq_len, num_sample]
- res = torch.ones_like(dtime_boundary_) * dtime_boundary_
+ res = self.sample_accept(unif_numbers, intensity_upper_bound, total_intensities, exp_numbers)
 
  # [batch_size, seq_len, num_sample]
- weights = torch.ones_like(dtime_boundary_)
- weights /= weights.sum(dim=-1, keepdim=True)
-
- res[who_has_accepted_times] = accepted_times_each_draw[who_has_accepted_times]
- who_not_accept = ~who_has_accepted_times
-
- who_reach_further = exp_numbers[..., -1] > dtime_boundary_
-
- res[who_not_accept & who_reach_further] = exp_numbers[..., -1][who_not_accept & who_reach_further]
-
+ weights = torch.ones_like(res)/res.shape[2]
+
  # add a upper bound here in case it explodes, e.g., in ODE models
  return res.clamp(max=1e5), weights