如何腌制从 TLorentzVector 继承的类？答案

【问题标题】：How to pickle classes which inherit from TLorentzVector?如何腌制从 TLorentzVector 继承的类？
【发布时间】：2014-09-08 22:20:58
【问题描述】：

这里使用 PyROOT (ROOT) 和酸洗。下面是一个非常简单的例子。我尝试以How to pickle an object of a class B (having many variables) that inherits from A, that defines __setstate__ and __getstate__ 为例。

from ROOT import TLorentzVector
import cPickle as pickle
class MyVec(TLorentzVector):
  def __init__(self):
    TLorentzVector.__init__(self)
    self.a = 'testing'

  def __getstate__(self):
    return self.__dict__

  def __setstate__(self, state):
    self.__dict__ = state

a = MyVec()
b = pickle.loads(pickle.dumps(a))

print a.__class__
print b.__class__

print a.__dict__
print b.__dict__

这个输出

<class '__main__.MyVec'>
<class 'ROOT.TLorentzVector'>
{'a': 'testing'}
{}

知道如何正确腌制我的对象吗？特别是，我不介意不从 TLorentzVector 继承并自己重载我正在使用的一些属性。但我仍然不清楚为什么我不能让它工作，因为它不保留对象属性。

【问题讨论】：

看来我需要在 MyVec 类上实现自定义的 reduce 方法，并且我需要正确执行。

标签： python python-2.7 pickle pyroot

【解决方案1】：

以下是我认为可行的示例。我不清楚它为什么起作用，但它一定与用原生 python 类包装 cython 类并为其定义方法有关......

#!/usr/bin/env python
from ROOT import TLorentzVector
import cPickle as pickle


class MyVec(TLorentzVector):
    def __init__(self, *args, **kwargs):
        print "init"
        if len(args) == 2 and isinstance(args[0], pow.__class__):
          super(MyVec, self).__init__(args[0](*args[1]))
        else:
          args = args or (TLorentzVector())
          if isinstance(args[0], TLorentzVector):
            super(MyVec, self).__init__(args[0])
          else:
            raise ValueError("Unexpected value")
        self.a = kwargs.get('a', 'fake')
        self.b = kwargs.get('b', TLorentzVector())
        print "args", args
        print "kwargs", kwargs

    def __setstate__(self, state):
        print "setstate"
        self.__dict__ = state

    def __getstate__(self):
        print "getstate"
        return self.__dict__

    def __reduce__(self):
        print "reduce"
        return (self.__class__, super(MyVec, self).__reduce__(), self.__getstate__(), )

anotherVec = TLorentzVector()
anotherVec.SetPtEtaPhiM(50.0, 0.0, 0.0, 0.0)

evenMore = TLorentzVector()
evenMore.SetPtEtaPhiM(100.0, 0.0, 0.0, 0.0)

a = MyVec(anotherVec, a='testing', b=evenMore)

b = pickle.loads(pickle.dumps(a))

print a.__class__
print b.__class__

print a.__dict__
print b.__dict__

print a.Pt()
print b.Pt()

print a.b.Pt()
print b.b.Pt()

print a == b
print isinstance(a, MyVec)
print isinstance(b, MyVec)

这个输出

init
args (<ROOT.TLorentzVector object ("TLorentzVector") at 0x7fa38bce26c0>,)
kwargs {'a': 'testing', 'b': <ROOT.TLorentzVector object ("TLorentzVector") at 0x7fa388f5c740>}
reduce
getstate
init
args (<built-in function _ObjectProxy__expand__>, ('@\x00\x00S\xff\xff\xff\xffTLorentzVector\x00@\x00\x00<\x00\x04\x00\x01\x00\x00\x00\x00\x03\x00\x00\x08@\x00\x00$\x00\x03\x00\x01\x00\x00\x00\x00\x03\x00\x00\x00@I\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00@I\x00\x00\x00\x00\x00\x00', 'TLorentzVector'))
kwargs {}
setstate
<class '__main__.MyVec'>
<class '__main__.MyVec'>
{'a': 'testing', 'b': <ROOT.TLorentzVector object ("TLorentzVector") at 0x7fa388f5c740>}
{'a': 'testing', 'b': <ROOT.TLorentzVector object ("TLorentzVector") at 0x7fa38b448be0>}
50.0
50.0
100.0
100.0
True
True
True

【讨论】：