fix bugs in template id mapping, and reduce memory usage

gstlal-inspiral/python/stats/inspiral_intrinsics.py

@@ -112,32 +112,28 @@ class SourcePopulationModel(object):
 		"""
 		if filename is not None:
 			with h5py.File(filename, 'r') as model:
-				coefficients = model['coefficients'][()]
 				snr_bp = model['SNR'][()]
 				try:
-					model_ids = model['template_id'][()]
+					template_indices = {x:n for n,x in enumerate(numpy.array(model['template_id']))}
 				except KeyError:
                                         # FIXME: assume sequential order if model['event_id'] doesn't exist
-                                        #model_ids = numpy.arange(numpy.shape(model['coefficients'][()])[-1])
-					model_ids = model['event_id'][()]
-			# PPoly can construct an array of polynomials by just
-			# feeding it the coefficients array all in one go, but then
-			# it insists on evaluating all of them at once.  we don't
-			# want to suffer that cost, so we have to make an array of
-			# PPoly objects ourselves, and index into it to evaluate
-			# just one.  since we have to do this anyway, we use a
-			# dictionary to also solve the problem of mapping
-			# template_id to a specific polynomial
-			template_indices = {}
-			for template_id in template_ids:
-                                # maps template ID to the right coefficient array, since template IDs
-                                # in the bank may not be in sequential order
+					template_indices = {x:n for n,x in enumerate(numpy.array(model['event_id']))}
+				# PPoly can construct an array of polynomials by just
+				# feeding it the coefficients array all in one go, but then
+				# it insists on evaluating all of them at once.  we don't
+				# want to suffer that cost, so we have to make an array of
+				# PPoly objects ourselves, and index into it to evaluate
+				# just one.  since we have to do this anyway, we use a
+				# dictionary to also solve the problem of mapping
+				# template_id to a specific polynomial
+				template_ids = sorted(template_ids)
 				try:
-					template_indices[template_id] = numpy.where(model_ids==template_id)[0][0]
-				except IndexError:
-					raise IndexError("template ID %d is not in this model" % template_id)
-			self.polys = dict((template_id, PPoly(coefficients[:,:,[template_indices[template_id]]], snr_bp)) for template_id in template_ids)
-			self.max_snr = snr_bp.max()
+					these_template_id_indices = numpy.array([template_indices[t] for t in template_ids], dtype=numpy.int)
+				except KeyError:
+					raise KeyError("One or more template IDs are not in this model")
+				coefficients = model['coefficients'][:,:,these_template_id_indices]
+				self.polys = dict((template_id, PPoly(coefficients[:,:, n], snr_bp)) for n, template_id in enumerate(template_ids))
+				self.max_snr = snr_bp.max()
 		else:
 			self.polys = None
 			self.max_snr = None
@@ -149,6 +145,5 @@ class SourcePopulationModel(object):
 			lnP_vs_snr = self.polys[template_id]
 		except KeyError:
 			raise KeyError("template ID %d is not in this model" % template_id)
-		# PPoly's .__call__() returns an array, so we need the
-		# final [0] to flatten it
-		return lnP_vs_snr(min(snr, self.max_snr))[0]
+		# PPoly's .__call__() returns an array, so we need the .item()
+		return lnP_vs_snr(min(snr, self.max_snr)).item()