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

Add a bit_depth attribute to G3Timestream objects, allowing the flac compression algorithm
to enable compression of the full 32 bits per sample, rather than restricting to 24 bits.

If 32-bit compression is selected, then int64 timestreams are truncated to int32, and
double timestreams are truncated to float prior to compression, and the resulting
data type is stored along with the compressed data.

Author: arahlin

core/CMakeLists.txt

@@ -94,6 +94,7 @@ add_spt3g_test(frame_types)
 add_spt3g_test(datatypes)
 if(FLAC_FOUND)
 	add_spt3g_test(ts_nanencoding)
+	add_spt3g_test(ts_flac_compression)
 endif()
 add_spt3g_test(cuts)
 add_spt3g_test(fileio)

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{
@@ -206,7 +211,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
@@ -225,9 +232,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.
@@ -243,16 +251,19 @@ class G3TimestreamMap : public G3FrameObject,
 	static G3TimestreamMap
 	MakeCompact(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) {
 		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) {
 		if(!std::is_sorted(keys.begin(), keys.end()))
 			throw std::runtime_error("G3TimestreamMap::MakeCompact: keys must be sorted");
 		const auto data_type=G3Timestream::TimeStreamTypeResolver<SampleType>::type_tag;
@@ -263,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;
@@ -283,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, 3);
 
 #endif

core/src/G3Timestream.cxx

@@ -122,9 +122,20 @@ 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;
+		if (flac_depth_ == 32) {
+			if (data_type_ == TS_INT64 || data_type_ == TS_INT32)
+				data_type_out = TS_INT32;
+		}
+
+		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_) {
+	        switch (flac_depth_) {
+		case 24:
+			switch (data_type_) {
 			case TS_DOUBLE:
 				for (size_t i = 0; i < size(); i++)
 					inbuf[i] = ((int32_t(((double *)data_)[i]) & 0x00ffffff) << 8) >> 8;
@@ -149,6 +160,31 @@ template <class A> void G3Timestream::save(A &ar, unsigned v) const
 				break;
 			default:
 				log_fatal("Unknown timestream datatype %d", data_type_);
+			}
+			break;
+		case 32:
+			switch (data_type_) {
+			case TS_DOUBLE:
+				for (size_t i = 0; i < size(); i++)
+					inbuf[i] = int32_t(float(((double *)data_)[i]));
+				break;
+			case TS_FLOAT:
+				for (size_t i = 0; i < size(); i++)
+					inbuf[i] = int32_t(((float *)data_)[i]);
+				break;
+			case TS_INT32:
+				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]);
+				break;
+			default:
+				log_fatal("Unknown timestream datatype %d", data_type_);
+			}
+			break;
+		default:
+			log_fatal("Invalid FLAC bit depth %d", flac_depth_);
 		}
 		chanmap[0] = &inbuf[0];
 
@@ -182,7 +218,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 +298,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 +324,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 +412,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 +482,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);
@@ -760,12 +819,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
@@ -1065,7 +1144,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);
 
@@ -1112,6 +1192,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) {
@@ -1343,6 +1424,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.")
@@ -1357,6 +1442,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__")
@@ -1378,7 +1465,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> "
@@ -1394,6 +1481,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")
@@ -1413,6 +1504,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>();
 

core/tests/G3TimestreamMapTest.cxx

@@ -55,10 +55,12 @@ TEST(MakeCompactCompressed){
 	std::vector<std::string> keys={"a","b","c","d"};
 	G3Time t1(0), t2(3);
 	const std::size_t nSamples=4;
-	G3TimestreamMap tsm=G3TimestreamMap::MakeCompact<int32_t>(keys, nSamples, t1, t2, G3Timestream::Counts, 6);
+	G3TimestreamMap tsm=G3TimestreamMap::MakeCompact<int32_t>(keys, nSamples, t1, t2, G3Timestream::Counts, 6, 32);
 
-	for(const auto& item : tsm)
+	for(const auto& item : tsm) {
 		ENSURE_EQUAL(item.second->GetCompressionLevel(), 6, "All timestreams should be set to use compression");
+		ENSURE_EQUAL(item.second->GetBitDepth(), 32, "All timestreams should be set to use 32-bit depth");
+	}
 }
 
 template<typename SampleType>
@@ -144,8 +146,11 @@ TEST(SetCompression){
 	G3TimestreamMap tsm=G3TimestreamMap::MakeCompact<double>(keys, nSamples, t1, t2);
 
 	unsigned int compressionLevel=5;
+	unsigned int bitDepth=32;
 	tsm.SetFLACCompression(5);
+	tsm.SetFLACDepth(bitDepth);
 	for(const auto& key : keys){
 		ENSURE_EQUAL(tsm[key]->GetCompressionLevel(), compressionLevel, "SetFLACCompression should set the compression level for each timestream");
+		ENSURE_EQUAL(tsm[key]->GetBitDepth(), bitDepth, "SetFLACDepth should set the bit depth for each timestream");
 	}
 }

core/tests/timestream_2d.py

@@ -64,9 +64,10 @@
 # Test setting secondary properties at the time of direct construction
 tsm2 = core.G3TimestreamMap(tsm.keys(), buffer1d, start=start, stop=stop,
                             units=core.G3TimestreamUnits.Counts,
-                            compression_level=3)
+                            compression_level=3, bit_depth=32)
 for ts in tsm2.values():
 	assert ts.start == start
 	assert ts.stop == stop
 	assert ts.units == core.G3TimestreamUnits.Counts
 	assert ts.compression_level == 3
+	assert ts.bit_depth == 32