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