happi: operation on scalars

doc/Sphinx/Overview/releases.rst

@@ -28,6 +28,8 @@ Changes made in the repository (not released)
   * New arguments ``xoffset`` and ``yoffset`` to shift plot coordinates.
   * Changed coordinate reference for 2D probe in 3D or AM geometry
     (zero is the box origin projected orthogonally on the probe plane).
+  * In ``Scalar``, it is now possible to make an operation on scalars such as ``"Uelm+Ukin"``.
+    The list of available scalars can be obtained from ``getScalars()``.
 
 * Bug fixes:
 

doc/Sphinx/Use/post-processing.rst

@@ -101,12 +101,16 @@ about the corresponding diagnostics in the simulation.
   Returns a list of available diagnostics of the given type
 
   * ``diagType``: The diagnostic type (``"Field"``, ``"Probe"``, etc.)
+
+.. rubric:: Information on specific diagnostics
+
+.. py:method:: getScalars()
+  
+  Returns a list of available scalars.
 
 .. py:method:: getTrackSpecies()
 
   Returns a list of available tracked species.
-
-.. rubric:: Information on specific diagnostics
 
 .. py:method:: fieldInfo(diag)
 
@@ -145,8 +149,7 @@ Open a Scalar diagnostic
 
 .. py:method:: Scalar(scalar=None, timesteps=None, units=[""], data_log=False, data_transform=None, **kwargs)
 
-  * ``scalar``: The name of the scalar.
-     | If not given, then a list of available scalars is printed.
+  * ``scalar``: The name of the scalar, or an operation on scalars, such as ``"Uelm+Ukin"``.
   * ``timesteps``: The requested timestep(s).
      | If omitted, all timesteps are used.
      | If one number  given, the nearest timestep available is used.

happi/_Diagnostics/Scalar.py

@@ -6,39 +6,55 @@ class Scalar(Diagnostic):
 
 	def _init(self, scalar=None, timesteps=None, data_log=False, data_transform=None, **kwargs):
 		# Get available scalars
-		scalars = self.getScalars()
+		scalars = self.simulation.getScalars()
 
 		# If no scalar chosen, only print the available scalars
 		if scalar is None:
+			error = []
 			if len(scalars)>0:
-				self._error += ["Error: no scalar chosen"]
-				self._error += ["Printing available scalars:"]
-				self._error += ["---------------------------"]
+				error += ["Error: no scalar chosen"]
+				error += ["Printing available scalars:"]
+				error += ["---------------------------"]
 				l = [""]
 				for s in scalars:
 					if len(s)>4 and s[:2]!=l[-1][:2] and s[-2:]!=l[-1][-2:]:
-						if l!=[""]: self._error += ["\t".join(l)]
+						if l!=[""]: error += ["\t".join(l)]
 						l = []
 					l.append(s)
-				if l!=[""]: self._error += ["\t".join(l)]
+				if l!=[""]: error += ["\t".join(l)]
 			else:
-				self._error += ["No scalars found"]
-			return
+				error += ["No scalars found"]
+			raise Exception("\n".join(error))
 
 		# 1 - verifications, initialization
 		# -------------------------------------------------------------------
-		# Find which scalar is requested
-		if scalar not in scalars:
-			fs = list(filter(lambda x:scalar in x, scalars))
-			if len(fs)==0:
-				self._error += ["No scalar `"+scalar+"` found"]
-				return
-			if len(fs)>1:
-				self._error += ["Several scalars match: "+(' '.join(fs))]
-				self._error += ["Please be more specific and retry."]
-				return
-			scalar = fs[0]
-		self._scalarname = scalar
+		# Find which scalar is requested & associated units
+		self.operation = scalar
+		def scalarTranslator(s):
+			units = "??"
+			if   s == "time":
+				units = "T_r"
+			elif s == "Ubal_norm":
+				units = ""
+			else:
+				units = {
+					"U":"K_r * N_r * L_r^%i" % self._ndim_particles,
+					"P":"K_r * N_r * L_r^%i" % self._ndim_particles,
+					"D":"N_r * L_r^%i" % self._ndim_particles,
+					"E":"E_r",
+					"B":"B_r",
+					"J":"J_r",
+					"R":"Q_r * N_r",
+					"Z":"Q_r",
+					"N":"",
+					}[s[0]]
+			return units, "S['%s']"%s, s
+		self._operation = Operation(scalar, scalarTranslator, self._ureg)
+		self._scalarname = self._operation.variables
+		self._vunits = self._operation.translated_units
+		self._title  = self._operation.title
+		if not self._scalarname:
+			raise Exception("String "+self.operation+" does not seem to include any scalar")
 
 		# Put data_log as object's variable
 		self._data_log = data_log
@@ -47,26 +63,31 @@ def _init(self, scalar=None, timesteps=None, data_log=False, data_transform=None
 		# Already get the data from the file
 		# Loop file line by line
 		self._alltimesteps = []
-		self._values = []
-		times_values = {}
+		S = { s:[] for s in self._scalarname }
 		for path in self._results_path:
 			try:
-				with open(path+'/scalars.txt') as f:
-					for line in f:
-						line = line.strip()
-						if line[0]!="#": break
-						prevline = line
-					scalars = prevline[1:].strip().split() # list of scalars
-					scalarindex = scalars.index(scalar) # index of the requested scalar
-					line = str(line.strip()).split()
-					times_values[ int( self._np.round(float(line[0]) / float(self.timestep)) ) ] = float(line[scalarindex])
-					for line in f:
-						line = str(line.strip()).split()
-						times_values[ int( self._np.round(float(line[0]) / float(self.timestep)) ) ] = float(line[scalarindex])
+				f = open(path+'/scalars.txt')
 			except:
 				continue
-		self._alltimesteps  = self._np.array(sorted(times_values.keys()))
-		self._values = self._np.array([times_values[k] for k in self._alltimesteps])
+			for line in f:
+				if line[0]!="#": break
+				prevline = line
+			scalars = prevline[1:].strip().split() # list of scalars
+			scalarindexes = {s:scalars.index(s) for s in self._scalarname} # indexes of the requested scalars
+			line = str(line.strip()).split()
+			self._alltimesteps += [ int( self._np.round(float(line[0]) / float(self.timestep)) ) ]
+			for s,i in scalarindexes.items():
+				S[s] += [float(line[i])]
+			for line in f:
+				line = str(line.strip()).split()
+				self._alltimesteps += [ int( self._np.round(float(line[0]) / float(self.timestep)) ) ]
+				for s,i in scalarindexes.items():
+					S[s] += [float(line[i])]
+			f.close()
+		self._alltimesteps  = self._np.array(self._alltimesteps)
+		for s in S:
+			S[s] = self._np.array(S[s])
+		self._values = self._operation.eval(locals())
 		self._timesteps = self._np.copy(self._alltimesteps)
 
 		# 2 - Manage timesteps
@@ -88,65 +109,13 @@ def _init(self, scalar=None, timesteps=None, data_log=False, data_transform=None
 			self._error += ["Timesteps not found"]
 			return
 
-
-		# 3 - Build units
-		# -------------------------------------------------------------------
-		self._vunits = "??"
-		if   self._scalarname == "time":
-			self._vunits = "T_r"
-		elif self._scalarname == "Ubal_norm":
-			self._vunits = ""
-		else:
-			self._vunits = {
-				"U":"K_r * N_r * L_r^%i" % self._ndim_particles,
-				"P":"K_r * N_r * L_r^%i" % self._ndim_particles,
-				"D":"N_r * L_r^%i" % self._ndim_particles,
-				"E":"E_r",
-				"B":"B_r",
-				"J":"J_r",
-				"R":"Q_r * N_r",
-				"Z":"Q_r",
-				"N":"",
-				}[self._scalarname[0]]
-		self._title =self._scalarname
-
 		# Finish constructor
 		self.valid = True
 		return kwargs
 
 	# Method to print info on included scalars
 	def _info(self):
-		return "Scalar "+self._scalarname
-
-	# get all available scalars
-	def getScalars(self):
-		allScalars = None
-		for path in self._results_path:
-			try:
-				file = path+'/scalars.txt'
-				f = open(file, 'r')
-			except Exception as e:
-				continue
-			try:
-				# Find last commented line
-				prevline = ""
-				for line in f:
-					line = line.strip()
-					if line[0]!="#": break
-					prevline = line[1:].strip()
-				scalars = str(prevline).split() # list of scalars
-				scalars = scalars[1:] # remove first, which is "time"
-			except Exception as e:
-				scalars = []
-			f.close()
-			if allScalars is None:
-				allScalars = scalars
-			else:
-				allScalars = self._np.intersect1d(allScalars, scalars)
-		if allScalars is None:
-			self._error += ["Cannot open 'scalars.txt'"]
-			return []
-		return allScalars
+		return "Scalar "+self.operation
 
 	# get all available timesteps
 	def getAvailableTimesteps(self):

happi/_Factories.py

@@ -39,10 +39,8 @@ def __init__(self, simulation, scalar=None):
 		# If not a specific scalar (root level), build a list of scalar shortcuts
 		if scalar is None:
 			if simulation._verbose: print("Scanning for Scalar diagnostics")
-			# Create a temporary, empty scalar diagnostic
-			tmpDiag = Scalar(simulation)
 			# Get a list of scalars
-			scalars = tmpDiag.getScalars()
+			scalars = simulation.getScalars()
 			# Create scalars shortcuts
 			for scalar in scalars:
 				setattr(self, scalar, ScalarFactory(simulation, scalar))
@@ -53,7 +51,25 @@ def __init__(self, simulation, scalar=None):
 
 	def __call__(self, *args, **kwargs):
 		return Scalar(self._simulation, *(self._additionalArgs+args), **kwargs)
-
+
+	def __repr__(self):
+		msg = object.__repr__(self)
+		if len(self._additionalArgs) == 0 and self._simulation._scan:
+			scalars = self._simulation.getScalars()
+			if scalars:
+				msg += "\nAvailable Scalar diagnostics:\n"
+				l = [""]
+				for s in scalars:
+					if len(s)>4 and s[:2]!=l[-1][:2] and s[-2:]!=l[-1][-2:]:
+						if l!=[""]:
+							msg += "\t".join(l) + "\n"
+						l = []
+					l.append(s)
+				if l!=[""]:
+					msg += "\t".join(l) + "\n"
+			else:
+				msg += "\nNo Scalar diagnostics available"
+		return msg
 
 class FieldFactory(object):
 	"""Import and analyze a Field diagnostic from a Smilei simulation

happi/_SmileiSimulation.py

@@ -281,6 +281,35 @@ def _getParticleListSpecies(self, filePrefix):
 				if s: species += [ s.groups()[0] ]
 		return list(set(species)) # unique species
 
+	# get all available scalars
+	def getScalars(self):
+		allScalars = None
+		for path in self._results_path:
+			try:
+				file = path+'/scalars.txt'
+				f = open(file, 'r')
+			except Exception as e:
+				continue
+			try:
+				# Find last commented line
+				prevline = ""
+				for line in f:
+					line = line.strip()
+					if line[0]!="#": break
+					prevline = line[1:].strip()
+				scalars = str(prevline).split() # list of scalars
+				scalars = scalars[1:] # remove first, which is "time"
+			except Exception as e:
+				scalars = []
+			f.close()
+			if allScalars is None:
+				allScalars = scalars
+			else:
+				allScalars = self._np.intersect1d(allScalars, scalars)
+		if allScalars is None:
+			return []
+		return allScalars
+
 	def getTrackSpecies(self):
 		""" List the available tracked species """
 		return self._getParticleListSpecies("TrackParticlesDisordered")

happi/_Utils.py

@@ -331,8 +331,7 @@ class Operation(object):
 	
 	Parameters:
 	-----------
-	operation: the user's requested operation
-	pattern: the regexp pattern to find the variables in the operation
+	operation: the user's requested operation (string containing variables and operators)
 	QuantityTranslator: function that takes a string as argument (a quantity name)
 		and outputs its units + its replacement string + its displayed name
 	ureg: Pint's unit registry or None for no unit awareness

validation/analyses/validate_tst1d_04_radiation_pressure_acc.py

@@ -31,6 +31,9 @@
 # SCALARS RELATED TO SPECIES
 Validate("Scalar Ukin_ion" , S.Scalar.Ukin_ion ().getData(), 100.)
 
+# OPERATION ON SCALARS
+Validate("Scalar Ukin+Uelm", S.Scalar("Ukin+Uelm").getData(), 100.)
+
 # 1D SCREEN DIAGS
 for i,d in enumerate(S.namelist.DiagScreen):
 	last_data = S.Screen(i, timesteps=21000).getData()[-1]