correldata

  1"""
  2Read/write vectors of correlated data from/to a csv file.
  3
  4These data are stored in a dictionary, whose values are numpy arrays
  5with elements which may be strings, floats, or floats with associated uncertainties
  6as defined in the [uncertainties](https://pypi.org/project/uncertainties) library.
  7"""
  8
  9
 10__author__    = 'Mathieu Daëron'
 11__contact__   = 'mathieu@daeron.fr'
 12__copyright__ = 'Copyright (c) 2024 Mathieu Daëron'
 13__license__   = 'MIT License - https://opensource.org/licenses/MIT'
 14__date__      = '2024-10-15'
 15__version__   = '1.3.0'
 16
 17
 18import os as _os
 19import numpy as _np
 20import uncertainties as _uc
 21
 22from typing import Callable, Hashable, Any
 23
 24class uarray(_np.ndarray):
 25
 26	__doc__ = """
 27	1-D [ndarray](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html)
 28	of [ufloat](https://pypi.org/project/uncertainties) values
 29	"""
 30
 31	def __new__(cls, a):
 32		obj = _np.asarray(a).view(cls)
 33		return obj
 34	
 35	@property
 36	def nv(self):
 37		"""Return the array of nominal values (read-only)."""
 38		return _uc.unumpy.nominal_values(_np.array(self))
 39
 40	@property
 41	def se(self):
 42		"""Return the array of standard errors (read-only)"""
 43		return _uc.unumpy.std_devs(_np.array(self))
 44
 45	@property
 46	def correl(self):
 47		"""Return the correlation matrix of the array elements (read-only)"""
 48		return _np.array(_uc.correlation_matrix(self))
 49
 50	@property
 51	def covar(self):
 52		"""Return the covariance matrix of the array elements (read-only)"""
 53		return _np.array(_uc.covariance_matrix(self))
 54	
 55	@property
 56	def mahalanobis(self):
 57		"""Return the squared Mahalanobis distance from zero of the array (read-only)"""
 58		flatself = self.n.flatten().reshape((1, self.size))
 59		return (flatself @ _np.linalg.inv(self.covar) @ flatself.T)[0,0]
 60	
 61	n = nv
 62	"Alias for `uarray.nv`"
 63	
 64	s = se
 65	"Alias for `uarray.se`"
 66	
 67	cor = correl
 68	"Alias for `uarray.correl`"
 69	
 70	cov = covar
 71	"Alias for `uarray.covar`"
 72	
 73	m = mahalanobis
 74	"Alias for `uarray.mahalanobis`"
 75
 76
 77def is_symmetric_positive_semidefinite(M: _np.ndarray) -> bool:
 78	'''
 79	Test whether 2-D array `M` is symmetric and positive semidefinite.
 80	'''
 81	ev = _np.linalg.eigvals(M)
 82	return (
 83		_np.allclose(M, M.T) # M is symmetric
 84		and _np.all(
 85			(ev > 0) | _np.isclose(ev, 0)
 86		) # all eignevalues are either real and strictly positive or close to zero
 87	)
 88
 89
 90def smart_type(s: str) -> (int | float | str):
 91	'''
 92	Tries to convert string `s` to an `int`, or to an `float` if that fails.
 93	If both fail, return the original string unchanged.
 94	'''
 95	try: return int(s)
 96	except: pass
 97	try: return float(s)
 98	except: pass
 99	return s
100
101
102def read_data(data: str, sep: str = ',', validate_covar: bool = True):
103	'''
104	Read correlated data from a CSV-like string.
105	
106	Column names are interpreted in the following way:
107	* In most cases, each columns is converted to a dict value, with the corresponding
108	dict key being the column's label.
109	* Columns whose label starts with `SE` are interpreted as specifying the standard
110	error for the latest preceding data column.
111	* Columns whose label starts with `correl` are interpreted as specifying the
112	correlation matrix for the latest preceding data column. In that case, column labels
113	are ignored for the rest of the columns belonging to this matrix.
114	* Columns whose label starts with `covar` are interpreted as specifying the
115	covariance matrix for the latest preceding data column. In that case, column labels
116	are ignored for the rest of the columns belonging to this matrix.
117	* `SE`, `correl`, and `covar` may be specified for any arbitrary variable other than
118	the latest preceding data column, by adding an underscore followed by the variable's
119	label (ex: `SE_foo`, `correl_bar`, `covar_baz`).
120	* `correl`, and `covar` may also be specified for any pair of variable, by adding an
121	underscore followed by the two variable labels, joined by a second underscore
122	(ex: `correl_foo_bar`, `covar_X_Y`). The elements of the first and second variables
123	correspond, respectively, to the lines and columns of this matrix.
124	* Exceptions will be raised, for any given variable:
125		- when specifying both `covar` and any combination of (`SE`, `correl`)
126		- when specifying `correl` without `SE`
127
128	**Arguments**
129	- `data`: a CSV-like string
130	- `sep`: the CSV separator
131	- `validate_covar`: whether to check that the overall covariance matrix
132	is symmetric and positive semidefinite. Specifying `validate_covar = False`
133	bypasses this computationally expensive step.
134	
135	**Example**
136	```py
137	import correldata
138	data  = """
139	Sample, Tacid,  D47,   SE,         correl,,,  D48, covar,,,          correl_D47_D48
140	   FOO,   90., .245, .005,      1, 0.5, 0.5, .145,  4e-4, 1e-4, 1e-4, 0.5,   0,   0
141	   BAR,   90., .246, .005,    0.5,   1, 0.5, .146,  1e-4, 4e-4, 1e-4,   0, 0.5,   0
142	   BAZ,   90., .247, .005,    0.5, 0.5,   1, .147,  1e-4, 1e-4, 4e-4,   0,   0, 0.5
143	"""[1:-1]
144	print(correldata.read_data(data))
145	
146	# yields:
147	# 
148	# > {
149	#     'Sample': array(['FOO', 'BAR', 'BAZ'], dtype='<U3'),
150	#     'Tacid': array([90., 90., 90.]),
151	#     'D47': uarray([0.245+/-0.004999999999999998, 0.246+/-0.004999999999999997, 0.247+/-0.005], dtype=object),
152	#     'D48': uarray([0.145+/-0.019999999999999993, 0.146+/-0.019999999999999993, 0.147+/-0.019999999999999997], dtype=object)
153	#   }
154	```
155	'''
156
157	data = [[smart_type(e.strip()) for e in l.split(sep)] for l in data.split('\n')]
158	N = len(data) - 1
159
160	values, se, correl, covar = {}, {}, {}, {}
161	j = 0
162	while j < len(data[0]):
163		field = data[0][j]
164		if not (
165			field.startswith('SE_')
166			or field.startswith('correl_')
167			or field.startswith('covar_')
168			or field == 'SE'
169			or field == 'correl'
170			or field == 'covar'
171			or len(field) == 0
172		):
173			values[field] = _np.array([l[j] for l in data[1:]])
174			j += 1
175			oldfield = field
176		elif field.startswith('SE_'):
177			se[field[3:]] = _np.array([l[j] for l in data[1:]])
178			j += 1
179		elif field == 'SE':
180			se[oldfield] = _np.array([l[j] for l in data[1:]])
181			j += 1
182		elif field.startswith('correl_'):
183			correl[field[7:]] = _np.array([l[j:j+N] for l in data[1:]])
184			j += N
185		elif field == 'correl':
186			correl[oldfield] = _np.array([l[j:j+N] for l in data[1:]])
187			j += N
188		elif field.startswith('covar_'):
189			covar[field[6:]] = _np.array([l[j:j+N] for l in data[1:]])
190			j += N
191		elif field == 'covar':
192			covar[oldfield] = _np.array([l[j:j+N] for l in data[1:]])
193			j += N
194
195	nakedvalues = {}
196	for k in [_ for _ in values]:
197		if (
198			k not in se
199			and k not in correl
200			and k not in covar
201		):
202			nakedvalues[k] = values.pop(k)
203
204	for x in values:
205		if x in covar:
206			if x in se:
207				raise KeyError(f'Too much information: both SE and covar are specified for variable "{x}".')
208			if x in correl:
209				raise KeyError(f'Too much information: both correl and covar are specified for variable "{x}".')
210		if x in correl:
211			if x not in se:
212				raise KeyError(f'Not enough information: correl is specified without SE for variable "{x}".')
213
214	for x in correl:
215		if x in values:
216			covar[x] = _np.diag(se[x]) @ correl[x] @ _np.diag(se[x])
217		else:
218			for x1 in values:
219				for x2 in values:
220					if x == f'{x1}_{x2}':
221						if x1 in se:
222							se1 = se[x1]
223						else:
224							if x1 in covar:
225								se1 = _np.diag(covar[x1])**0.5
226							else:
227								raise KeyError(f'Not enough information: correl_{x} is specified without SE for variable "{x1}".')
228						if x2 in se:
229							se2 = se[x2]
230						else:
231							if x2 in covar:
232								se2 = _np.diag(covar[x2])**0.5
233							else:
234								raise KeyError(f'Not enough information: correl_{x} is specified without SE for variable "{x1}".')
235
236						covar[x] = _np.diag(se1) @ correl[x] @ _np.diag(se2)
237
238	for x in se:
239		if x in values and x not in correl:
240			covar[x] = _np.diag(se[x]**2)
241
242	for k in [_ for _ in covar]:
243		if k not in values:
244			for j1 in values:
245				for j2 in values:
246					if k == f'{j1}_{j2}':
247						covar[f'{j2}_{j1}'] = covar[f'{j1}_{j2}'].T
248
249	X = _np.array([_ for k in values for _ in values[k]])
250	CM = _np.zeros((X.size, X.size))
251	for i, vi in enumerate(values):
252		for j, vj in enumerate(values):
253			if vi == vj:
254				if vi in covar:
255					CM[N*i:N*i+N,N*j:N*j+N] = covar[vi]
256			else:
257				if f'{vi}_{vj}' in covar:
258					CM[N*i:N*i+N,N*j:N*j+N] = covar[f'{vi}_{vj}']
259
260	s = _np.diag(CM)**.5
261	s[s==0] = 1.
262	invs = _np.diag(s**-1)
263
264	if (
265		validate_covar
266		and not (
267			is_symmetric_positive_semidefinite(CM)
268			or is_symmetric_positive_semidefinite(invs @ CM @ invs)
269		)
270	):
271		raise _np.linalg.LinAlgError('The complete covariance matrix is not symmetric positive-semidefinite.')
272
273	corvalues = uarray(_uc.correlated_values(X, CM))
274
275	allvalues = nakedvalues
276
277	for i, x in enumerate(values):
278		allvalues[x] = corvalues[i*N:i*N+N]
279
280	return allvalues
281
282
283def read_data_from_file(filename: str | _os.PathLike, **kwargs):
284	'''
285	Read correlated data from a CSV file.
286
287	**Arguments**
288	- `filename`: `str` or path to the file to read from
289	- `kwargs`: passed to correldata.read_data()
290	'''
291	with open(filename) as fid:
292		return read_data(fid.read(), **kwargs)
293
294
295def f2s(
296	x: Any,
297	f: (str | Callable | dict),
298	k: Hashable = None,
299	fb: (str | Callable) = 'z.6g',
300) -> str:
301	'''
302	Format `x` according to format `f`
303	
304	* If `f` is a string, return `f'{x:{f}}'`
305	* If `f` is a callable, return `f(x)`
306	* If `f` is a dict and optional argument `k` is a hashable,
307	  return f2s(x, f[k]), otherwise return f2s(x, fb)
308	'''
309
310	if isinstance (x, str):
311		return x
312	if isinstance (f, str):
313		return f'{x:{f}}'
314	if isinstance (f, Callable):
315		return f(x)
316	if isinstance (f, dict):
317		if k in f:
318			return f2s(x, f[k])
319		if isinstance (fb, str):
320			return f'{x:{fb}}'
321		if isinstance (fb, Callable):
322			return fb(x)
323	raise TypeError(f'f2s() formatting argument f = {repr(f)} is neither a string nor a dict nor a callable.')
324	
325
326
327def data_string(
328	data: dict,
329	sep: str = ',',
330	include_fields: list = None,
331	exclude_fields: list = [],
332	float_format: (str | dict | Callable) = 'z.6g',
333	correl_format: (str | dict | Callable) = 'z.6f',
334	default_float_format: (str | Callable) = 'z.6g',
335	default_correl_format: (str | Callable) = 'z.6f',
336	align: str = '>',
337	atol: float = 1e-12,
338	rtol: float = 1e-12,
339):
340	'''
341	Generate CSV-like string from correlated data
342
343	**Arguments**
344	- `data`: dict of arrays with strings, floats or correlated data
345	- `sep`: the CSV separator
346	- `include_fields`: subset of fields to write; if `None`, write all fields
347	- `exclude_fields`: subset of fields to ignore (takes precedence over `include_fields`);
348	  to exclude only the SE for field `foo`, include `SE_foo`; same goes for `correl_foo`
349	- `float_format`: formatting for float values. May be a string (ex: `'z.3f'`), a callable
350	  (ex: `lambda x: '.2f' if x else '0'`), or a dictionary of strings and/or callables, with dict keys
351	  corresponding to different fields (ex: `{'foo': '.2e', 'bar': (lambda x: str(x))}`).
352	- `correl_format`: same as `float_format`, but applies to correlation matrix elements
353	- `default_float_format`: only used when `float_format` is a dict; in that case, fields
354	  missing from `float_format.keys()` will use `default_float_format` instead.
355	  corresponding to different fields (ex: `{'foo': '.2e', 'bar': `lambda x: str(x)`}`).
356	- `default_correl_format`: same as `default_float_format`, but applies to `correl_format`
357	- `align`: right-align (`>`), left-align (`<`), or don't align (empty string) CSV values
358	- `atol`: passed to [numpy.allclose()](https://numpy.org/doc/stable/reference/generated/numpy.allclose.html)
359	  when deciding whether a matrix is equal to the identity matrix or to the zero matrix
360	- `rtol`: passed to [numpy.allclose()](https://numpy.org/doc/stable/reference/generated/numpy.allclose.html)
361	  when deciding whether a matrix is equal to the identity matrix or to the zero matrix
362	
363	
364	**Example**
365	
366	```py
367	from correldata import _uc
368	from correldata import _np
369	from correldata import *
370	
371	X = uarray(_uc.correlated_values([1., 2., 3.], _np.eye(3)*0.09))
372	Y = uarray(_uc.correlated_values([4., 5., 6.], _np.eye(3)*0.16))
373	
374	data = dict(X=X, Y=Y, Z=X+Y)
375	
376	print(data_string(data, float_format = 'z.1f', correl_format = 'z.1f'))
377	
378	# yields:
379	# 
380	#   X, SE_X,   Y, SE_Y,   Z, SE_Z, correl_X_Z,    ,    , correl_Y_Z,    ,    
381	# 1.0,  0.3, 4.0,  0.4, 5.0,  0.5,        0.6, 0.0, 0.0,        0.8, 0.0, 0.0
382	# 2.0,  0.3, 5.0,  0.4, 7.0,  0.5,        0.0, 0.6, 0.0,        0.0, 0.8, 0.0
383	# 3.0,  0.3, 6.0,  0.4, 9.0,  0.5,        0.0, 0.0, 0.6,        0.0, 0.0, 0.8
384	```
385	'''
386	if include_fields is None:
387		include_fields = [_ for _ in data]
388	cols, ufields = [], []
389	for f in include_fields:
390		if f in exclude_fields:
391			continue
392		if isinstance(data[f], uarray):
393			ufields.append(f)
394			N = data[f].size
395			cols.append([f] + [f2s(_, float_format, f, default_float_format) for _ in data[f].n])
396			if f'SE_{f}' not in exclude_fields:
397				cols.append([f'SE_{f}'] + [f2s(_, float_format, f, default_float_format) for _ in data[f].s])
398			if f'correl_{f}' not in exclude_fields:
399				CM = _uc.correlation_matrix(data[f])
400				if not _np.allclose(CM, _np.eye(N), atol = atol, rtol = rtol):
401					for i in range(N):
402						cols.append(
403							['' if i else f'correl_{f}']
404							+ [
405								f2s(
406									CM[i,j],
407									correl_format,
408									f,
409									default_correl_format,
410								)
411								for j in range(N)
412							]
413						)
414
415		else:
416			cols.append([f] + [f2s(_, float_format, f, default_float_format) for _ in data[f]])
417
418	for i in range(len(ufields)):
419		for j in range(i):
420			if f'correl_{ufields[i]}_{ufields[j]}' in exclude_fields or f'correl_{ufields[j]}_{ufields[i]}' in exclude_fields:
421				continue
422			CM = _uc.correlation_matrix((*data[ufields[i]], *data[ufields[j]]))[:N, -N:]
423			if not _np.allclose(CM, _np.zeros((N, N)), atol = atol, rtol = rtol):
424				for k in range(N):
425					cols.append(
426						['' if k else f'correl_{ufields[j]}_{ufields[i]}']
427						+ [
428							f2s(
429								CM[k,l],
430								correl_format,
431								f,
432								default_correl_format,
433							)
434							for l in range(N)
435						]
436					)
437
438	lines = list(map(list, zip(*cols)))
439
440	if align:
441		lengths = [max([len(e) for e in l]) for l in cols]
442		for l in lines:
443			for k,ln in enumerate(lengths):
444				l[k] = f'{l[k]:{align}{ln}s}'
445		return '\n'.join([(sep+' ').join(l) for l in lines])
446
447	return '\n'.join([sep.join(l) for l in lines])
448
449
450
451def save_data_to_file(data, filename, **kwargs):
452	'''
453	Write correlated data to a CSV file.
454
455	**Arguments**
456	- `data`: dict of arrays with strings, floats or correlated data
457	- `filename`: `str` or path to the file to read from
458	- `kwargs`: passed to correldata.data_string()
459	'''
460	with open(filename, 'w') as fid:
461		return fid.write(data_string(data, **kwargs))