simulate_experiment with domain error when n_outcomes is not constant

[PengDauan]

I was following the examples file：introduction_to_bayes_smc, I find the following issue when I was trying to call the simulate_experiment methods

[PengDauan]

class FiniteOutcomeModel(Model):

def domain(self, expparams):
    if self.is_n_outcomes_constant:
        return self._domain if expparams is None else [self._domain for ep in expparams]
    else:
        return [IntegerDomain(min=0,max=n_o-1) for n_o in self.n_outcomes(expparams)]

def simulate_experiment(self, modelparams, expparams, repeat=1):

    super(FiniteOutcomeModel, self).simulate_experiment(modelparams, expparams, repeat)
    
    if self.is_n_outcomes_constant:
        all_outcomes = self.domain(None).values
        probabilities = self.likelihood(all_outcomes, modelparams, expparams)
        cdf = np.cumsum(probabilities, axis=0)
        randnum = np.random.random((repeat, 1, modelparams.shape[0], expparams.shape[0]))
        outcome_idxs = all_outcomes[np.argmax(cdf > randnum, axis=1)]
        outcomes = all_outcomes[outcome_idxs]
    else:
        assert(self.are_expparam_dtypes_consistent(expparams))
        dtype = self.domain(expparams[0, np.newaxis])[0].dtype
        outcomes = np.empty((repeat, modelparams.shape[0], expparams.shape[0]), dtype=dtype)
        for idx_experiment, single_expparams in enumerate(expparams[:, np.newaxis]):
            all_outcomes = self.domain(single_expparams).values
            probabilities = self.likelihood(all_outcomes, modelparams, single_expparams)
            cdf = np.cumsum(probabilities, axis=0)[..., 0]
            randnum = np.random.random((repeat, 1, modelparams.shape[0]))
            outcomes[:, :, idx_experiment] = all_outcomes[np.argmax(cdf > randnum, axis=1)]
            
    return outcomes[0, 0, 0] if repeat == 1 and expparams.shape[0] == 1 and modelparams.shape[0] == 1 else outcomes

[PengDauan]

def simulate_experiment(self, modelparams, expparams, repeat=1):
	# Call the superclass simulate_experiment, not recording the result.
	# This is used to count simulation calls.
	super(FiniteOutcomeModel, self).simulate_experiment(modelparams, expparams, repeat)
	
	if self.is_n_outcomes_constant:
		# In this case, all expparams have the same domain
		all_outcomes = self.domain(None).values
		probabilities = self.likelihood(all_outcomes, modelparams, expparams)
		cdf = np.cumsum(probabilities, axis=0)
		randnum = np.random.random((repeat, 1, modelparams.shape[0], expparams.shape[0]))
		outcome_idxs = all_outcomes[np.argmax(cdf > randnum, axis=1)]
		outcomes = all_outcomes[outcome_idxs]
	else:
		# Loop over each experiment, sadly.
		# Assume all domains have the same dtype
		assert(self.are_expparam_dtypes_consistent(expparams))
		dtype = self.domain(expparams[0, np.newaxis])[0].dtype
		outcomes = np.empty((repeat, modelparams.shape[0], expparams.shape[0]), dtype=dtype)
		for idx_experiment, single_expparams in enumerate(expparams[:, np.newaxis]):
			all_outcomes = self.domain(single_expparams).values
			probabilities = self.likelihood(all_outcomes, modelparams, single_expparams)
			cdf = np.cumsum(probabilities, axis=0)[..., 0]
			randnum = np.random.random((repeat, 1, modelparams.shape[0]))
			outcomes[:, :, idx_experiment] = all_outcomes[np.argmax(cdf > randnum, axis=1)]
			
	return outcomes[0, 0, 0] if repeat == 1 and expparams.shape[0] == 1 and modelparams.shape[0] == 1 else outcomes

change the code in abstract.FiniteOutcomeModel.simulate_experiment
all_outcomes= self.domain(single_expparams).values
to
all_outcomes= self.domain(single_expparams)[idx_experiment].values
will solve the problem in this case, but I'm not sure this is the general case