support passing arrays to FieldDesc + CategoryDesc and CategoryWriter as @DataClass

dsehnal · dsehnal · commit 02b81952cdc9 · 2022-09-30T13:48:56.000+02:00
diff --git a/ciftools/binary/writer.py b/ciftools/binary/writer.py
@@ -101,18 +101,32 @@ def _encode_field(field: FieldDesc, data: list[_ContextData], total_count: int)
 offset = 0
 for _d in data:
 d = _d.data
- for i in range(_d.count):
- p = presence(d, i)
- if p is not ValuePresenceEnum.Present:
- mask[offset] = p
- if is_native:
- array[offset] = None
- all_present = False
- else:
- mask[offset] = ValuePresenceEnum.Present
- array[offset] = field.value(d, i)
-
- offset += 1
+
+ arrays = field.arrays(d)
+ if arrays is not None:
+ if len(arrays.values) != _d.count:
+ raise ValueError(f"values provided in arrays() must have the same length as the category count field")
+
+ if arrays.mask is not None:
+ if len(arrays.mask) != _d.count:
+ raise ValueError(f"mask provided in arrays() must have the same length as the category count field")
+ mask[offset:offset + _d.count] = arrays.mask
+ offset += _d.count
+
+ else:
+ # TODO: use numba JIT for this
+ for i in range(_d.count):
+ p = presence(d, i)
+ if p is not ValuePresenceEnum.Present:
+ mask[offset] = p
+ if is_native:
+ array[offset] = None
+ all_present = False
+ else:
+ mask[offset] = ValuePresenceEnum.Present
+ array[offset] = field.value(d, i)
+
+ offset += 1
 
 encoder = field.encoder(data[0].data) if len(data) > 0 else _BYTE_ARRAY_ENCODER
 encoded = encoder.encode_cif_data(array)
diff --git a/ciftools/writer/base.py b/ciftools/writer/base.py
@@ -1,10 +1,16 @@
 import abc
-from typing import Any, Union
+from dataclasses import dataclass
+from typing import Any, List, Optional, Union
 
 import numpy as np
 from ciftools.binary.encoding.impl.binary_cif_encoder import BinaryCIFEncoder
 from ciftools.cif_format.value_presence import ValuePresenceEnum
 
+@dataclass
+class FieldArrays:
+ values: Union[np.ndarray, List[str]]
+ # uint8 array, 0 = defined, 1 = ., 2 = ?
+ mask: Optional[np.ndarray] = None
 
 class FieldDesc(abc.ABC):
 name: str
@@ -17,6 +23,10 @@ def create_array(self, total_count: int) -> Union[np.ndarray, list]:
 def value(self, data: Any, i: int) -> Any:
 pass
 
+ @abc.abstractmethod
+ def arrays(self, data: Any) -> Optional[FieldArrays]:
+ pass
+
 @abc.abstractmethod
 def encoder(self, data: Any) -> BinaryCIFEncoder:
 pass
@@ -25,12 +35,13 @@ def encoder(self, data: Any) -> BinaryCIFEncoder:
 def presence(self, data: any, i: int) -> ValuePresenceEnum:
 pass
 
-
+@dataclass
 class CategoryDesc:
 name: str
 fields: list[FieldDesc]
 
 
+@dataclass
 class CategoryWriter:
 data: any
 count: int
diff --git a/ciftools/writer/fields.py b/ciftools/writer/fields.py
@@ -4,7 +4,7 @@
 from ciftools.binary.encoding.impl.binary_cif_encoder import BinaryCIFEncoder
 from ciftools.binary.encoding.impl.encoders.string_array import STRING_ARRAY_CIF_ENCODER
 from ciftools.cif_format.value_presence import ValuePresenceEnum
-from ciftools.writer.base import FieldDesc
+from ciftools.writer.base import FieldArrays, FieldDesc
 
 _STRING_ARRAY_ENCODER = BinaryCIFEncoder([STRING_ARRAY_CIF_ENCODER])
 
@@ -23,24 +23,30 @@ def create_array(self, total_count: int):
 def encoder(self, data: Any):
 return _STRING_ARRAY_ENCODER
 
+ def arrays(self, data: Any) -> Optional[FieldArrays]:
+ return self._arrays(data) if self._arrays is not None else None
+
 def __init__(
 self,
 name: str,
 value: Callable[[Any, int], Optional[str]],
 presence: Optional[Callable[[Any, int], Optional[ValuePresenceEnum]]] = None,
+ arrays: Optional[Callable[[Any], FieldArrays]] = None,
 ) -> None:
 self.name = name
 self._value = value
 self._presence = presence
+ self._arrays = arrays
 
 
 def string_field(
 *,
 name: str,
 value: Callable[[Any, int], Optional[str]],
 presence: Optional[Callable[[Any, int], Optional[ValuePresenceEnum]]] = None,
+ arrays: Optional[Callable[[Any], FieldArrays]] = None,
 ) -> FieldDesc:
- return _StringFieldDesc(name=name, value=value, presence=presence)
+ return _StringFieldDesc(name=name, value=value, presence=presence, arrays=arrays)
 
 
 # TODO: derive from FieldDesc
@@ -57,27 +63,33 @@ def presence(self, data: any, i: int) -> ValuePresenceEnum:
 def create_array(self, total_count: int):
 return np.empty(total_count, dtype=self._dtype)
 
+ def arrays(self, data: Any) -> Optional[FieldArrays]:
+ return self._arrays(data) if self._arrays is not None else None
+
 def __init__(
 self,
 name: str,
 value: Callable[[Any, int], Optional[Union[int, float]]],
 dtype: np.dtype,
 encoder: Callable[[Any], BinaryCIFEncoder],
 presence: Optional[Callable[[Any, int], Optional[ValuePresenceEnum]]] = None,
+ arrays: Optional[Callable[[Any], FieldArrays]] = None
 ) -> None:
 self.name = name
 self._value = value
 self._dtype = dtype
 self._encoder = encoder
 self._presence = presence
+ self._arrays = arrays
 
 
 def number_field(
 *,
 name: str,
- value: Callable[[Any, int], Optional[Union[int, float]]],
+ value: Optional[Callable[[Any, int], Optional[Union[int, float]]]] = None,
 dtype: np.dtype,
 encoder: Callable[[Any], BinaryCIFEncoder],
 presence: Optional[Callable[[Any, int], Optional[ValuePresenceEnum]]] = None,
+ arrays: Optional[Callable[[Any], FieldArrays]] = None,
 ) -> FieldDesc:
- return _NumberFieldDesc(name=name, value=value, dtype=dtype, encoder=encoder, presence=presence)
+ return _NumberFieldDesc(name=name, value=value, dtype=dtype, encoder=encoder, presence=presence, arrays=arrays)
diff --git a/tests/_encoding.py b/tests/_encoding.py
@@ -8,7 +8,7 @@
 from ciftools.binary.encoding.impl.encoders.integer_packing import INTEGER_PACKING_CIF_ENCODER
 from ciftools.binary.writer import BinaryCIFWriter
 from ciftools.cif_format.binary.file import BinaryCIFFile
-from ciftools.writer.base import CategoryDesc, CategoryWriter, CategoryWriterProvider, FieldDesc, OutputStream
+from ciftools.writer.base import CategoryDesc, CategoryWriter, CategoryWriterProvider, FieldArrays, FieldDesc, OutputStream
 from ciftools.writer.fields import number_field, string_field
 
 
@@ -102,6 +102,17 @@ def lattice_value_getter(lid: int):
 value=lambda data, i: data.volume[i],
 )
 )
+ fields.append(
+ number_field(
+ name=f"volume_array",
+ dtype="f4",
+ encoder=lambda _: BinaryCIFEncoder(
+ [FixedPointCIFEncoder(1000), DELTA_CIF_ENCODER, INTEGER_PACKING_CIF_ENCODER]
+ ),
+ value=lambda data, i: data.volume[i],
+ arrays=lambda data: FieldArrays(values=data.volume),
+ )
+ )
 fields.append(string_field(name="annotation", value=lambda data, i: data.annotation[i]))
 
 return TestCategoryWriter(ctx, self.length, TestCategoryDesc("volume", fields))
@@ -169,6 +180,11 @@ def test(self):
 compare = np.allclose(test_data.volume, volume, atol=1e-3)
 self.assertTrue(compare, "Volume did not match original data")
 
+ volume_array = volume_and_lattices.get_column("volume_array").__dict__["_values"]
+ print("Volume Array (parsed): " + str(volume_array))
+ compare = np.allclose(test_data.volume, volume_array, atol=1e-3)
+ self.assertTrue(compare, "Volume Array did not match original data")
+
 for lattice_id in lattice_ids:
 print("Lattice: " + str(lattice_id))
 lattice_value = volume_and_lattices.get_column("lattice_" + str(lattice_id)).__dict__["_values"]

-Original file line number
+Diff line change
 offset=0
 for_dindata:
 d=_d.data
 -foriinrange(_d.count):
 -p=presence(d, i)
 -ifpisnotValuePresenceEnum.Present:
 -mask[offset] =p
 -ifis_native:
 -array[offset] =None
 -all_present=False
 -else:
 -mask[offset] =ValuePresenceEnum.Present
 -array[offset] =field.value(d, i)
+-
 -offset+=1
++
 +arrays=field.arrays(d)
 +ifarraysisnotNone:
 +iflen(arrays.values) !=_d.count:
 +raiseValueError(f"values provided in arrays() must have the same length as the category count field")
++
 +ifarrays.maskisnotNone:
 +iflen(arrays.mask) !=_d.count:
 +raiseValueError(f"mask provided in arrays() must have the same length as the category count field")
 +mask[offset:offset+_d.count] =arrays.mask
 +offset+=_d.count
++
 +else:
 +# TODO: use numba JIT for this
 +foriinrange(_d.count):
 +p=presence(d, i)
 +ifpisnotValuePresenceEnum.Present:
 +mask[offset] =p
 +ifis_native:
 +array[offset] =None
 +all_present=False
 +else:
 +mask[offset] =ValuePresenceEnum.Present
 +array[offset] =field.value(d, i)
++
 +offset+=1
 encoder=field.encoder(data[0].data) iflen(data) >0else_BYTE_ARRAY_ENCODER
 encoded=encoder.encode_cif_data(array)
-Original file line number
+Diff line change
@@ @@ -1,10 +1,16 @@ @@
 importabc
 -fromtypingimportAny, Union
 +fromdataclassesimportdataclass
 +fromtypingimportAny, List, Optional, Union
 importnumpyasnp
 fromciftools.binary.encoding.impl.binary_cif_encoderimportBinaryCIFEncoder
 fromciftools.cif_format.value_presenceimportValuePresenceEnum
 +@dataclass
 +classFieldArrays:
 +values: Union[np.ndarray, List[str]]
 +# uint8 array, 0 = defined, 1 = ., 2 = ?
 +mask: Optional[np.ndarray] =None
 classFieldDesc(abc.ABC):
 name: str
 defvalue(self, data: Any, i: int) ->Any:
 pass
 +@abc.abstractmethod
 +defarrays(self, data: Any) ->Optional[FieldArrays]:
 +pass
++
 @abc.abstractmethod
 defencoder(self, data: Any) ->BinaryCIFEncoder:
 pass
 defpresence(self, data: any, i: int) ->ValuePresenceEnum:
 pass
+-
 +@dataclass
 classCategoryDesc:
 name: str
 fields: list[FieldDesc]
 +@dataclass
 classCategoryWriter:
 data: any
 count: int
-Original file line number
+Diff line change
 fromciftools.binary.encoding.impl.binary_cif_encoderimportBinaryCIFEncoder
 fromciftools.binary.encoding.impl.encoders.string_arrayimportSTRING_ARRAY_CIF_ENCODER
 fromciftools.cif_format.value_presenceimportValuePresenceEnum
 -fromciftools.writer.baseimportFieldDesc
 +fromciftools.writer.baseimportFieldArrays, FieldDesc
 _STRING_ARRAY_ENCODER=BinaryCIFEncoder([STRING_ARRAY_CIF_ENCODER])
 defencoder(self, data: Any):
 return_STRING_ARRAY_ENCODER
 +defarrays(self, data: Any) ->Optional[FieldArrays]:
 +returnself._arrays(data) ifself._arraysisnotNoneelseNone
++
 def__init__(
 self,
 name: str,
 value: Callable[[Any, int], Optional[str]],
 presence: Optional[Callable[[Any, int], Optional[ValuePresenceEnum]]] =None,
 +arrays: Optional[Callable[[Any], FieldArrays]] =None,
  ) ->None:
 self.name=name
 self._value=value
 self._presence=presence
 +self._arrays=arrays
 defstring_field(
 *,
 name: str,
 value: Callable[[Any, int], Optional[str]],
 presence: Optional[Callable[[Any, int], Optional[ValuePresenceEnum]]] =None,
 +arrays: Optional[Callable[[Any], FieldArrays]] =None,
 ) ->FieldDesc:
 -return_StringFieldDesc(name=name, value=value, presence=presence)
 +return_StringFieldDesc(name=name, value=value, presence=presence, arrays=arrays)
 # TODO: derive from FieldDesc
 defcreate_array(self, total_count: int):
 returnnp.empty(total_count, dtype=self._dtype)
 +defarrays(self, data: Any) ->Optional[FieldArrays]:
 +returnself._arrays(data) ifself._arraysisnotNoneelseNone
++
 def__init__(
 self,
 name: str,
 value: Callable[[Any, int], Optional[Union[int, float]]],
 dtype: np.dtype,
 encoder: Callable[[Any], BinaryCIFEncoder],
 presence: Optional[Callable[[Any, int], Optional[ValuePresenceEnum]]] =None,
 +arrays: Optional[Callable[[Any], FieldArrays]] =None
  ) ->None:
 self.name=name
 self._value=value
 self._dtype=dtype
 self._encoder=encoder
 self._presence=presence
 +self._arrays=arrays
 defnumber_field(
 *,
 name: str,
 -value: Callable[[Any, int], Optional[Union[int, float]]],
 +value: Optional[Callable[[Any, int], Optional[Union[int, float]]]] =None,
 dtype: np.dtype,
 encoder: Callable[[Any], BinaryCIFEncoder],
 presence: Optional[Callable[[Any, int], Optional[ValuePresenceEnum]]] =None,
 +arrays: Optional[Callable[[Any], FieldArrays]] =None,
 ) ->FieldDesc:
 -return_NumberFieldDesc(name=name, value=value, dtype=dtype, encoder=encoder, presence=presence)
 +return_NumberFieldDesc(name=name, value=value, dtype=dtype, encoder=encoder, presence=presence, arrays=arrays)
-Original file line number
+Diff line change
 fromciftools.binary.encoding.impl.encoders.integer_packingimportINTEGER_PACKING_CIF_ENCODER
 fromciftools.binary.writerimportBinaryCIFWriter
 fromciftools.cif_format.binary.fileimportBinaryCIFFile
 -fromciftools.writer.baseimportCategoryDesc, CategoryWriter, CategoryWriterProvider, FieldDesc, OutputStream
 +fromciftools.writer.baseimportCategoryDesc, CategoryWriter, CategoryWriterProvider, FieldArrays, FieldDesc, OutputStream
 fromciftools.writer.fieldsimportnumber_field, string_field
 value=lambdadata, i: data.volume[i],
+ )
+ )
 +fields.append(
 +number_field(
 +name=f"volume_array",
 +dtype="f4",
 +encoder=lambda_: BinaryCIFEncoder(
 + [FixedPointCIFEncoder(1000), DELTA_CIF_ENCODER, INTEGER_PACKING_CIF_ENCODER]
 + ),
 +value=lambdadata, i: data.volume[i],
 +arrays=lambdadata: FieldArrays(values=data.volume),
 + )
 + )
 fields.append(string_field(name="annotation", value=lambdadata, i: data.annotation[i]))
 returnTestCategoryWriter(ctx, self.length, TestCategoryDesc("volume", fields))
 compare=np.allclose(test_data.volume, volume, atol=1e-3)
 self.assertTrue(compare, "Volume did not match original data")
 +volume_array=volume_and_lattices.get_column("volume_array").__dict__["_values"]
 +print("Volume Array (parsed): "+str(volume_array))
 +compare=np.allclose(test_data.volume, volume_array, atol=1e-3)
 +self.assertTrue(compare, "Volume Array did not match original data")
++
 forlattice_idinlattice_ids:
 print("Lattice: "+str(lattice_id))
 lattice_value=volume_and_lattices.get_column("lattice_"+str(lattice_id)).__dict__["_values"]