Add support for full-depth compression and decompression of float and int32 timestreams

arahlin · arahlin · commit 18798c0dbdd3 · 2025-04-03T12:53:38.000-05:00
diff --git a/core/include/core/G3Timestream.h b/core/include/core/G3Timestream.h
@@ -27,19 +27,21 @@ class G3Timestream : public G3FrameObject {
 
 	G3Timestream(const G3Timestream &r);
 	G3Timestream(std::vector<double>::size_type s = 0, double val = 0) :
-	    units(None), use_flac_(0),
+	    units(None), use_flac_(0), flac_depth_(24),
 	    buffer_((s == 0) ? NULL : new std::vector<double>(s, val)),
 	    data_((s == 0) ? NULL : &(*buffer_)[0]), len_(s),
 	    data_type_(TS_DOUBLE) {}
 	template <typename Iterator> G3Timestream(Iterator l, Iterator r) :
-	    units(None), use_flac_(0), buffer_(new std::vector<double>(l, r)),
+	    units(None), use_flac_(0), flac_depth_(24),
+	    buffer_(new std::vector<double>(l, r)),
 	    data_(&(*buffer_)[0]), len_(buffer_->size()), data_type_(TS_DOUBLE) {}
 	virtual ~G3Timestream() {
 		if (buffer_) delete buffer_;
 	}
 
 	// FLAC compression levels range from 0-9. 0 means do not use FLAC.
 	void SetFLACCompression(int compression_level);
+	void SetFLACDepth(int bit_depth);
 
 	TimestreamUnits units;
 	G3Time start, stop;
@@ -105,6 +107,7 @@ class G3Timestream : public G3FrameObject {
 
 	double GetSampleRate() const;
 	uint8_t GetCompressionLevel() const{ return use_flac_; }
+	uint8_t GetBitDepth() const { return flac_depth_; }
 	
 	template <class A> void load(A &ar, unsigned v);
 	template <class A> void save(A &ar, unsigned v) const;
@@ -135,17 +138,19 @@ class G3Timestream : public G3FrameObject {
 	friend class G3TimestreamPythonHelpers;
 
 	uint8_t use_flac_;
+	uint8_t flac_depth_;
 
 	std::vector<double> *buffer_;
 	std::shared_ptr<void> root_data_ref_;
 	void *data_;
 	size_t len_;
-	enum {
+	enum DataType {
 		TS_DOUBLE,
 		TS_FLOAT,
 		TS_INT32,
 		TS_INT64
-	} data_type_;
+	};
+	DataType data_type_;
 	
 	template<typename T>
 	struct TimeStreamTypeResolver{
@@ -208,7 +213,9 @@ class G3TimestreamMap : public G3FrameObject,
 	void SetUnits(G3Timestream::TimestreamUnits units);
 	/// FLAC compression levels range from 0-9. 0 means do not use FLAC.
 	uint8_t GetCompressionLevel() const;
+	uint8_t GetBitDepth() const;
 	void SetFLACCompression(int compression_level);
+	void SetFLACDepth(int bit_depth);
 
 	// Compact underlying data storage into a contiguous 2D block.
 	// This invalidates any references to data inside any member
@@ -227,9 +234,10 @@ class G3TimestreamMap : public G3FrameObject,
 	static G3TimestreamMap MakeCompact(const std::vector<std::string>& keys, std::size_t n_samples,
 	                                   G3Time start, G3Time stop,
 	                                   G3Timestream::TimestreamUnits units=G3Timestream::None,
-	                                   int compression_level=0){
+	                                   int compression_level=0, int bit_depth=24){
 		std::shared_ptr<SampleType[]> data(new SampleType[n_samples*keys.size()]);
-		return MakeCompact(keys, n_samples, data, start, stop, units, compression_level);
+		return MakeCompact(keys, n_samples, data, start, stop, units,
+		    compression_level, bit_depth);
 	}
 
 	/// Construct a map using an existing contiguous 2D block of data as the underlying storage.
@@ -246,16 +254,18 @@ class G3TimestreamMap : public G3FrameObject,
 	                                   std::shared_ptr<SampleType[]> data,
 	                                   G3Time start, G3Time stop,
 	                                   G3Timestream::TimestreamUnits units=G3Timestream::None,
-	                                   int compression_level=0){
+	                                   int compression_level=0, int bit_depth=24){
 		G3TimestreamMap map;
-		map.FromBuffer(keys, n_samples, data, start, stop, units, compression_level);
+		map.FromBuffer(keys, n_samples, data, start, stop, units,
+		    compression_level, bit_depth);
 		return map;
 	}
 
 	template<typename SampleType>
 	void FromBuffer(const std::vector<std::string>& keys, std::size_t n_samples,
 	    std::shared_ptr<SampleType[]> data, G3Time start, G3Time stop,
-	    G3Timestream::TimestreamUnits units=G3Timestream::None, int compression_level=0) {
+	    G3Timestream::TimestreamUnits units=G3Timestream::None, int compression_level=0,
+	    int bit_depth=24) {
 		const auto data_type=G3Timestream::TimeStreamTypeResolver<SampleType>::type_tag;
 		std::size_t offset = 0;
 		for (const auto& key : keys) {
@@ -264,6 +274,7 @@ class G3TimestreamMap : public G3FrameObject,
 			ts->stop = stop;
 			ts->units = units;
 			ts->use_flac_ = compression_level;
+			ts->flac_depth_ = bit_depth;
 			ts->root_data_ref_ = data;
 			ts->data_ = data.get() + offset;
 			ts->data_type_ = data_type;
@@ -284,7 +295,7 @@ namespace cereal {
 	template <class A> struct specialize<A, G3TimestreamMap, cereal::specialization::member_serialize> {};
 }
 
-G3_SERIALIZABLE(G3Timestream, 3);
+G3_SERIALIZABLE(G3Timestream, 4);
 G3_SERIALIZABLE(G3TimestreamMap, 4);
 
 #endif
diff --git a/core/src/G3SuperTimestream.cxx b/core/src/G3SuperTimestream.cxx
@@ -248,13 +248,19 @@ template <class A> void G3SuperTimestream::load(A &ar, unsigned v)
 		FromBuffer(names, nsamp_filled, data, start, stop);
 		buf = (char *)data.get();
 		elsize = sizeof(int32_t);
+
+		SetFLACCompression(flac_level);
+		SetFLACDepth(32);
 		break;
 	}
 	case TYPE_NUM_FLOAT32: {
 		std::shared_ptr<float[]> data(new float[size]);
 		FromBuffer(names, nsamp_filled, data, start, stop);
 		buf = (char *)data.get();
 		elsize = sizeof(int32_t);
+
+		SetFLACCompression(flac_level);
+		SetFLACDepth(32);
 		break;
 	}
 	case TYPE_NUM_INT64: {
diff --git a/core/src/G3Timestream.cxx b/core/src/G3Timestream.cxx
@@ -102,6 +102,11 @@ extern "C"{
 }
 #endif
 
+template <typename T>
+inline int32_t truncate(void *data, size_t i, int32_t bitmask, int32_t bitshift) {
+	return ((int32_t(((T *)data)[i]) & bitmask) << bitshift) >> bitshift;
+}
+
 template <class A> void G3Timestream::save(A &ar, unsigned v) const
 {
 	ar & cereal::make_nvp("G3FrameObject",
@@ -122,30 +127,52 @@ template <class A> void G3Timestream::save(A &ar, unsigned v) const
 		if (units != Counts && units != None)
 			log_fatal("Cannot use FLAC on non-counts timestreams");
 
+		DataType data_type_out = TS_FLOAT;
+		int32_t bitmask, bitshift;
+		switch (flac_depth_) {
+		case 24:
+			bitmask = 0x00ffffff;
+			bitshift = 8;
+			break;
+		case 32:
+			bitmask = 0xffffffff;
+			bitshift = 0;
+			if (data_type_ == TS_INT64 || data_type_ == TS_INT32)
+				data_type_out = TS_INT32;
+			break;
+		default:
+			log_fatal("Invalid FLAC bit depth %d", (int) flac_depth_);
+		}
+
+		ar & cereal::make_nvp("flac_depth", flac_depth_);
+		ar & cereal::make_nvp("data_type", data_type_out);
+
 		// Copy to 24-bit integers
 		inbuf.resize(size());
 		switch (data_type_) {
 			case TS_DOUBLE:
 				for (size_t i = 0; i < size(); i++)
-					inbuf[i] = ((int32_t(((double *)data_)[i]) & 0x00ffffff) << 8) >> 8;
+					inbuf[i] = truncate<double>(data_, i, bitmask, bitshift);
 				break;
 			case TS_FLOAT:
 				for (size_t i = 0; i < size(); i++)
-					inbuf[i] = ((int32_t(((float *)data_)[i]) & 0x00ffffff) << 8) >> 8;
+					inbuf[i] = truncate<float>(data_, i, bitmask, bitshift);
 				break;
 			case TS_INT32:
-				{
+				if (flac_depth_ == 24) {
 					// Using this rather raw form for the loop can enable automatic
 					// unrolling and vectorization.
 					int32_t* in_ptr=(int32_t *)data_;
 					int32_t* out_ptr=&inbuf[0];
 					for(int32_t* end=in_ptr+size(); in_ptr!=end; in_ptr++,out_ptr++)
 						*out_ptr = ((*in_ptr & 0x00ffffff) << 8) >> 8;
+				} else {
+					memcpy(inbuf.data(), data_, size() * sizeof(int32_t));
 				}
 				break;
 			case TS_INT64:
 				for (size_t i = 0; i < size(); i++)
-					inbuf[i] = ((int32_t(((int64_t *)data_)[i]) & 0x00ffffff) << 8) >> 8;
+					inbuf[i] = truncate<int64_t>(data_, i, bitmask, bitshift);
 				break;
 			default:
 				log_fatal("Unknown timestream datatype %d", data_type_);
@@ -182,7 +209,7 @@ template <class A> void G3Timestream::save(A &ar, unsigned v) const
 		FLAC__StreamEncoder *encoder = FLAC__stream_encoder_new();
 		FLAC__stream_encoder_set_channels(encoder, 1);
 		// XXX: should assert if high-order 8 bits are not clear
-		FLAC__stream_encoder_set_bits_per_sample(encoder, 24);
+		FLAC__stream_encoder_set_bits_per_sample(encoder, flac_depth_);
 		FLAC__stream_encoder_set_compression_level(encoder, use_flac_);
 		FLAC__stream_encoder_set_do_md5(encoder, false);
 		FLAC__stream_encoder_init_stream(encoder,
@@ -262,6 +289,13 @@ template <class A> void G3Timestream::load(A &ar, unsigned v)
 		if (units != Counts && units != None)
 			log_fatal("Cannot use FLAC on non-counts timestreams");
 
+		if (v >= 4) {
+			ar & cereal::make_nvp("flac_depth", flac_depth_);
+			ar & cereal::make_nvp("data_type", data_type_);
+		} else {
+			data_type_ = TS_FLOAT;
+		}
+
 		ar & cereal::make_nvp("nanflag", nanflag);
 		if (nanflag == SomeNan)
 			ar & cereal::make_nvp("nanmask", nanbuf);
@@ -281,14 +315,22 @@ template <class A> void G3Timestream::load(A &ar, unsigned v)
 		FLAC__stream_decoder_finish(decoder);
 		FLAC__stream_decoder_delete(decoder);
 
+		// Short-circuit for full-depth INT32.
+		if (data_type_ == TS_INT32 && flac_depth_ == 32) {
+			root_data_ref_ = std::shared_ptr<std::vector<int32_t> >(
+			    callback.outbuf);
+			len_ = callback.outbuf->size();
+			data_ = &(*callback.outbuf)[0];
+			return;
+		}
+
 		// Represent data as floats internally. These have the same
 		// significand depth (24 bits) as the max. bit depth of the
 		// reference FLAC encoder we use, so no data are lost, and
 		// allow NaNs, unlike int32_ts, which we try to pull through
 		// the process to signal missing data.
 		float *data = new float[callback.outbuf->size()];
 		root_data_ref_ = std::shared_ptr<float[]>(data);
-		data_type_ = TS_FLOAT;
 		len_ = callback.outbuf->size();
 		data_ = data;
 
@@ -361,7 +403,7 @@ template <class A> void G3Timestream::load(A &ar, unsigned v)
 
 G3Timestream::G3Timestream(const G3Timestream &r) :
     units(r.units), start(r.start), stop(r.stop), use_flac_(r.use_flac_),
-    len_(r.len_), data_type_(r.data_type_)
+    flac_depth_(r.flac_depth_), len_(r.len_), data_type_(r.data_type_)
 {
 	// Copy constructor needs to copy data, which always involves
 	// allocating the internal buffer.
@@ -431,6 +473,14 @@ void G3Timestream::SetFLACCompression(int use_flac)
 #endif
 }
 
+void G3Timestream::SetFLACDepth(int bit_depth)
+{
+	if (bit_depth != 24 && bit_depth != 32)
+		log_fatal("Invalid flac bit depth %d", bit_depth);
+
+	flac_depth_ = bit_depth;
+}
+
 G3Timestream G3Timestream::operator +(const G3Timestream &r) const
 {
 	G3Timestream ret(*this);
@@ -764,12 +814,32 @@ uint8_t G3TimestreamMap::GetCompressionLevel() const
 	return begin()->second->use_flac_;
 }
 
+uint8_t G3TimestreamMap::GetBitDepth() const
+{
+	if (begin() == end())
+		return 0;
+
+	return begin()->second->flac_depth_;
+}
+
 void G3TimestreamMap::SetFLACCompression(int compression_level)
 {
+	// Check for errors
+	begin()->second->SetFLACCompression(compression_level);
+
 	for (auto& ts : *this)
 		ts.second->use_flac_ = compression_level;
 }
 
+void G3TimestreamMap::SetFLACDepth(int bit_depth)
+{
+	// Check for errors
+	begin()->second->SetFLACDepth(bit_depth);
+
+	for (auto& ts : *this)
+		ts.second->flac_depth_ = bit_depth;
+}
+
 void G3TimestreamMap::Compactify()
 {
 	// Check if already compacted
@@ -1069,7 +1139,8 @@ struct PyBufferOwner {
 static G3TimestreamMapPtr
 G3TimestreamMap_from_numpy(std::vector<std::string> keys,
     boost::python::object data, G3Time start, G3Time stop,
-    G3Timestream::TimestreamUnits units, int compression_level, bool copy_data)
+    G3Timestream::TimestreamUnits units, int compression_level,
+    bool copy_data, int bit_depth)
 {
 	G3TimestreamMapPtr x(new G3TimestreamMap);
 
@@ -1116,6 +1187,7 @@ G3TimestreamMap_from_numpy(std::vector<std::string> keys,
 	templ.start = start;
 	templ.stop = stop;
 	templ.SetFLACCompression(compression_level);
+	templ.SetFLACDepth(bit_depth);
 	if (strcmp(v->v.format, "d") == 0) {
 		templ.data_type_ = G3Timestream::TS_DOUBLE;
 	} else if (strcmp(v->v.format, "f") == 0) {
@@ -1347,6 +1419,10 @@ PYBINDINGS("core") {
 	      "Pass True to turn on FLAC compression when serialized. "
 	      "FLAC compression only works if the timestream is in units of "
 	      "counts.")
+	    .def("SetFLACDepth", &G3Timestream::SetFLACDepth,
+	      "Change the bit depth for FLAC compression. "
+	      "FLAC compression only works if the timestream is in units of "
+	      "counts.")
 	    .def_readwrite("units", &G3Timestream::units,
 	      "Units of the data in the timestream, stored as one of the "
 	      "members of core.G3TimestreamUnits.")
@@ -1361,6 +1437,8 @@ PYBINDINGS("core") {
 	    .add_property("compression_level", &G3Timestream::GetCompressionLevel,
 	      &G3Timestream::SetFLACCompression, "Level of FLAC compression used for this timestream. "
 	      "This can only be non-zero if the timestream is in units of counts.")
+	    .add_property("bit_depth", &G3Timestream::GetBitDepth,
+	      &G3Timestream::SetFLACDepth, "Bit depth of FLAC compression used for this timestream.")
 	    .def("_assert_congruence", &G3Timestream::G3TimestreamPythonHelpers::G3Timestream_assert_congruence,
 	      "log_fatal() if units, length, start, or stop times do not match")
 	    .def("_cxxslice", &G3Timestream::G3TimestreamPythonHelpers::G3Timestream_getslice, "Slice-only __getitem__")
@@ -1382,7 +1460,7 @@ PYBINDINGS("core") {
 	         bp::default_call_policies(),
 	         (bp::arg("keys"), bp::arg("data"), bp::arg("start")=G3Time(0),
 	          bp::arg("stop")=G3Time(0), bp::arg("units") = G3Timestream::TimestreamUnits::None,
-	          bp::arg("compression_level") = 0, bp::arg("copy_data") = true)),
+	          bp::arg("compression_level") = 0, bp::arg("copy_data") = true, bp::arg("bit_depth") = 24)),
 	         "Create a timestream map from a numpy array or other numeric python iterable. "
 	         "Each row of the 2D input array will correspond to a single timestream, with "
 	         "the key set to the correspondingly-indexed entry of <keys>. If <copy_data> "
@@ -1398,6 +1476,10 @@ PYBINDINGS("core") {
 	      "Pass True to turn on FLAC compression when serialized. "
 	      "FLAC compression only works if the timestreams are in units of "
 	      "counts.")
+	    .def("SetFLACDepth", &G3TimestreamMap::SetFLACDepth,
+	      "Change the bit depth for FLAC compression. "
+	      "FLAC compression only works if the timestreams are in units of "
+	      "counts.")
 	    .add_property("start", &G3TimestreamMap::GetStartTime,
 	      &G3TimestreamMap::SetStartTime,
 	      "Time of the first sample in the time stream")
@@ -1417,6 +1499,9 @@ PYBINDINGS("core") {
 	      &G3TimestreamMap::SetFLACCompression,
 	      "Level of FLAC compression used for this timestream map. "
 	      "This can only be non-zero if the timestream is in units of counts.")
+	    .add_property("bit_depth", &G3TimestreamMap::GetBitDepth,
+	      &G3TimestreamMap::SetFLACDepth,
+	      "Bit depth of FLAC compression used for this timestream map.")
 	;
 	register_pointer_conversions<G3TimestreamMap>();
 
