Commit 8c5da49356cdcfdbb41efc74a35616a40603424c

Authored by Jürgen Knödlseder
1 parent 91643221

Support Python list slicing in container classes (#2128)

pyext/GEnergies.i
... ... @@ -79,17 +79,55 @@ public:
79 79 ***************************************************************************/
80 80 %extend GEnergies {
81 81 GEnergy& __getitem__(const int& index) {
  82 + // Counting from start, e.g. [2]
82 83 if (index >= 0 && index < self->size()) {
83 84 return (*self)[index];
84 85 }
  86 + // Counting from end, e.g. [-1]
  87 + else if (index < 0 && self->size()+index >= 0) {
  88 + return (*self)[self->size()+index];
  89 + }
85 90 else {
86 91 throw GException::out_of_range("__getitem__(int)", index, self->size());
87 92 }
88 93 }
  94 + GEnergies* __getitem__(PyObject *param) {
  95 + if (PySlice_Check(param)) {
  96 + Py_ssize_t start = 0;
  97 + Py_ssize_t stop = 0;
  98 + Py_ssize_t step = 0;
  99 + Py_ssize_t len = self->size();
  100 + if (PySlice_GetIndices((PySliceObject*)param, len, &start, &stop, &step) == 0) {
  101 + GEnergies* energies = new GEnergies;
  102 + if (step > 0) {
  103 + for (int i = (int)start; i < (int)stop; i += (int)step) {
  104 + energies->append((*self)[i]);
  105 + }
  106 + }
  107 + else {
  108 + for (int i = (int)start; i > (int)stop; i += (int)step) {
  109 + energies->append((*self)[i]);
  110 + }
  111 + }
  112 + return energies;
  113 + }
  114 + else {
  115 + throw GException::invalid_argument("__getitem__(PyObject)",
  116 + "Invalid slice indices");
  117 + }
  118 + }
  119 + else {
  120 + throw GException::invalid_argument("__getitem__(PyObject)","");
  121 + }
  122 + }
89 123 void __setitem__(const int& index, const GEnergy& val) {
90   - if (index>=0 && index < self->size()) {
  124 + // Counting from start, e.g. [2]
  125 + if (index >= 0 && index < self->size()) {
91 126 (*self)[index] = val;
92   - return;
  127 + }
  128 + // Counting from end, e.g. [-1]
  129 + else if (index < 0 && self->size()+index >= 0) {
  130 + (*self)[self->size()+index] = val;
93 131 }
94 132 else {
95 133 throw GException::out_of_range("__setitem__(int)", index, self->size());
... ...
pyext/GPhotons.i
... ... @@ -60,17 +60,55 @@ public:
60 60 ***************************************************************************/
61 61 %extend GPhotons {
62 62 GPhoton& __getitem__(const int& index) {
  63 + // Counting from start, e.g. [2]
63 64 if (index >= 0 && index < self->size()) {
64 65 return (*self)[index];
65 66 }
  67 + // Counting from end, e.g. [-1]
  68 + else if (index < 0 && self->size()+index >= 0) {
  69 + return (*self)[self->size()+index];
  70 + }
66 71 else {
67 72 throw GException::out_of_range("__getitem__(int)", index, self->size());
68 73 }
69 74 }
  75 + GPhotons* __getitem__(PyObject *param) {
  76 + if (PySlice_Check(param)) {
  77 + Py_ssize_t start = 0;
  78 + Py_ssize_t stop = 0;
  79 + Py_ssize_t step = 0;
  80 + Py_ssize_t len = self->size();
  81 + if (PySlice_GetIndices((PySliceObject*)param, len, &start, &stop, &step) == 0) {
  82 + GPhotons* photons = new GPhotons;
  83 + if (step > 0) {
  84 + for (int i = (int)start; i < (int)stop; i += (int)step) {
  85 + photons->append((*self)[i]);
  86 + }
  87 + }
  88 + else {
  89 + for (int i = (int)start; i > (int)stop; i += (int)step) {
  90 + photons->append((*self)[i]);
  91 + }
  92 + }
  93 + return photons;
  94 + }
  95 + else {
  96 + throw GException::invalid_argument("__getitem__(PyObject)",
  97 + "Invalid slice indices");
  98 + }
  99 + }
  100 + else {
  101 + throw GException::invalid_argument("__getitem__(PyObject)","");
  102 + }
  103 + }
70 104 void __setitem__(const int& index, const GPhoton& val) {
71   - if (index>=0 && index < self->size()) {
  105 + // Counting from start, e.g. [2]
  106 + if (index >= 0 && index < self->size()) {
72 107 (*self)[index] = val;
73   - return;
  108 + }
  109 + // Counting from end, e.g. [-1]
  110 + else if (index < 0 && self->size()+index >= 0) {
  111 + (*self)[self->size()+index] = val;
74 112 }
75 113 else {
76 114 throw GException::out_of_range("__setitem__(int)", index, self->size());
... ...
pyext/GTimes.i
1 1 /***************************************************************************
2 2 * GTimes.i - Time container class *
3 3 * ----------------------------------------------------------------------- *
4   - * copyright (C) 2012-2015 by Juergen Knoedlseder *
  4 + * copyright (C) 2012-2017 by Juergen Knoedlseder *
5 5 * ----------------------------------------------------------------------- *
6 6 * *
7 7 * This program is free software: you can redistribute it and/or modify *
... ... @@ -62,17 +62,55 @@ public:
62 62 ***************************************************************************/
63 63 %extend GTimes {
64 64 GTime& __getitem__(const int& index) {
  65 + // Counting from start, e.g. [2]
65 66 if (index >= 0 && index < self->size()) {
66 67 return (*self)[index];
67 68 }
  69 + // Counting from end, e.g. [-1]
  70 + else if (index < 0 && self->size()+index >= 0) {
  71 + return (*self)[self->size()+index];
  72 + }
68 73 else {
69 74 throw GException::out_of_range("__getitem__(int)", index, self->size());
70 75 }
71 76 }
  77 + GTimes* __getitem__(PyObject *param) {
  78 + if (PySlice_Check(param)) {
  79 + Py_ssize_t start = 0;
  80 + Py_ssize_t stop = 0;
  81 + Py_ssize_t step = 0;
  82 + Py_ssize_t len = self->size();
  83 + if (PySlice_GetIndices((PySliceObject*)param, len, &start, &stop, &step) == 0) {
  84 + GTimes* times = new GTimes;
  85 + if (step > 0) {
  86 + for (int i = (int)start; i < (int)stop; i += (int)step) {
  87 + times->append((*self)[i]);
  88 + }
  89 + }
  90 + else {
  91 + for (int i = (int)start; i > (int)stop; i += (int)step) {
  92 + times->append((*self)[i]);
  93 + }
  94 + }
  95 + return times;
  96 + }
  97 + else {
  98 + throw GException::invalid_argument("__getitem__(PyObject)",
  99 + "Invalid slice indices");
  100 + }
  101 + }
  102 + else {
  103 + throw GException::invalid_argument("__getitem__(PyObject)","");
  104 + }
  105 + }
72 106 void __setitem__(const int& index, const GTime& val) {
73   - if (index>=0 && index < self->size()) {
  107 + // Counting from start, e.g. [2]
  108 + if (index >= 0 && index < self->size()) {
74 109 (*self)[index] = val;
75   - return;
  110 + }
  111 + // Counting from end, e.g. [-1]
  112 + else if (index < 0 && self->size()+index >= 0) {
  113 + (*self)[self->size()+index] = val;
76 114 }
77 115 else {
78 116 throw GException::out_of_range("__setitem__(int)", index, self->size());
... ...
test/test_GObservation.py
... ... @@ -70,6 +70,46 @@ class Test(gammalib.GPythonTestSuite):
70 70 # Return observations container
71 71 return obs
72 72  
  73 + # Setup GTimes container
  74 + def _setup_times(self):
  75 + """
  76 + Setup GTimes container
  77 +
  78 + Returns
  79 + -------
  80 + times : `~gammalib.GTimes`
  81 + GTime container
  82 + """
  83 + # Setup time container
  84 + times = gammalib.GTimes()
  85 + time = gammalib.GTime()
  86 + for i in range(10):
  87 + time.secs(float(i))
  88 + times.append(time)
  89 +
  90 + # Return times container
  91 + return times
  92 +
  93 + # Setup GPhotons container
  94 + def _setup_photons(self):
  95 + """
  96 + Setup GPhotons container
  97 +
  98 + Returns
  99 + -------
  100 + photons : `~gammalib.GPhotons`
  101 + GPhoton container
  102 + """
  103 + # Setup photon container
  104 + photons = gammalib.GPhotons()
  105 + photon = gammalib.GPhoton()
  106 + for i in range(10):
  107 + photon.energy().MeV(float(i))
  108 + photons.append(photon)
  109 +
  110 + # Return photons container
  111 + return photons
  112 +
73 113 # Test GEnergy class
74 114 def _test_energy(self):
75 115 """
... ... @@ -149,9 +189,11 @@ class Test(gammalib.GPythonTestSuite):
149 189 """
150 190 Test GEnergies class
151 191 """
152   - # Setup 3 logarithmic energies
  192 + # Setup energies container
153 193 energies = gammalib.GEnergies(3, gammalib.GEnergy(1.0, 'MeV'),
154 194 gammalib.GEnergy(100.0, 'MeV'))
  195 +
  196 + # Check energies container
155 197 self.test_value(energies[0].MeV(), 1.0, 1.0e-7,
156 198 'Check energy value access')
157 199 self.test_value(energies[1].MeV(), 10.0, 1.0e-7,
... ... @@ -162,6 +204,91 @@ class Test(gammalib.GPythonTestSuite):
162 204 # Return
163 205 return
164 206  
  207 + # Test GEnergies class access operators
  208 + def _test_energies_access(self):
  209 + """
  210 + Test GEnergies class observation access
  211 + """
  212 + # Setup energies container
  213 + energies = gammalib.GEnergies(10, gammalib.GEnergy(1.0, 'MeV'),
  214 + gammalib.GEnergy(10.0, 'MeV'), False)
  215 +
  216 + # Loop over all elements using the container iterator and count the
  217 + # number of iterations
  218 + nenergies = 0
  219 + reference = 1.0
  220 + for energy in energies:
  221 + self.test_value(energy.MeV(), reference)
  222 + nenergies += 1
  223 + reference += 1.0
  224 +
  225 + # Check that looping was successful
  226 + self.test_value(energies.size(), 10, 'Check container size')
  227 + self.test_value(nenergies, 10, 'Check container iterator')
  228 +
  229 + # Test access from start
  230 + self.test_value(energies[3].MeV(), 4.0)
  231 +
  232 + # Check access from end
  233 + self.test_value(energies[-2].MeV(), 9.0)
  234 +
  235 + # Return
  236 + return
  237 +
  238 + # Test GEnergies class slicing
  239 + def _test_energies_slicing(self):
  240 + """
  241 + Test GEnergies class slicing
  242 + """
  243 + # Setup energies container
  244 + energies = gammalib.GEnergies(10, gammalib.GEnergy(1.0, 'MeV'),
  245 + gammalib.GEnergy(10.0, 'MeV'), False)
  246 +
  247 + # Test energies[start:end]
  248 + self.test_value(len(energies[3:5]), 2)
  249 + self.test_value(energies[3:5][0].MeV(), 4.0)
  250 + self.test_value(energies[3:5][1].MeV(), 5.0)
  251 +
  252 + # Test energies[start:]
  253 + self.test_value(len(energies[7:]), 3)
  254 + self.test_value(energies[7:][0].MeV(), 8.0)
  255 + self.test_value(energies[7:][1].MeV(), 9.0)
  256 + self.test_value(energies[7:][2].MeV(), 10.0)
  257 +
  258 + # Test energies[:end]
  259 + self.test_value(len(energies[:2]), 2)
  260 + self.test_value(energies[:2][0].MeV(), 1.0)
  261 + self.test_value(energies[:2][1].MeV(), 2.0)
  262 +
  263 + # Test energies[:]
  264 + self.test_value(len(energies[:]), 10)
  265 + for i in range(10):
  266 + self.test_value(energies[:][i].MeV(), float(i+1))
  267 +
  268 + # Test energies[start:end:step]
  269 + self.test_value(len(energies[3:7:2]), 2)
  270 + self.test_value(energies[3:7:2][0].MeV(), 4.0)
  271 + self.test_value(energies[3:7:2][1].MeV(), 6.0)
  272 +
  273 + # Test energies[start:end:step]
  274 + self.test_value(len(energies[6:3:-2]), 2)
  275 + self.test_value(energies[6:3:-2][0].MeV(), 7.0)
  276 + self.test_value(energies[6:3:-2][1].MeV(), 5.0)
  277 +
  278 + # Test energies[-start:]
  279 + self.test_value(len(energies[-2:]), 2)
  280 + self.test_value(energies[-2:][0].MeV(), 9.0)
  281 + self.test_value(energies[-2:][1].MeV(), 10.0)
  282 +
  283 + # Test energies[:-end]
  284 + self.test_value(len(energies[:-7]), 3)
  285 + self.test_value(energies[:-7][0].MeV(), 1.0)
  286 + self.test_value(energies[:-7][1].MeV(), 2.0)
  287 + self.test_value(energies[:-7][2].MeV(), 3.0)
  288 +
  289 + # Return
  290 + return
  291 +
165 292 # Test GTime class
166 293 def _test_time(self):
167 294 """
... ... @@ -231,6 +358,172 @@ class Test(gammalib.GPythonTestSuite):
231 358 # Return
232 359 return
233 360  
  361 + # Test GTimes class access operators
  362 + def _test_times_access(self):
  363 + """
  364 + Test GTimes class observation access
  365 + """
  366 + # Setup times container
  367 + times = self._setup_times()
  368 +
  369 + # Loop over all elements using the container iterator and count the
  370 + # number of iterations
  371 + ntimes = 0
  372 + reference = 0.0
  373 + for time in times:
  374 + self.test_value(time.secs(), reference)
  375 + ntimes += 1
  376 + reference += 1.0
  377 +
  378 + # Check that looping was successful
  379 + self.test_value(times.size(), 10, 'Check container size')
  380 + self.test_value(ntimes, 10, 'Check container iterator')
  381 +
  382 + # Test access from start
  383 + self.test_value(times[3].secs(), 3.0)
  384 +
  385 + # Check access from end
  386 + self.test_value(times[-2].secs(), 8.0)
  387 +
  388 + # Return
  389 + return
  390 +
  391 + # Test GTimes class slicing
  392 + def _test_times_slicing(self):
  393 + """
  394 + Test GTimes class slicing
  395 + """
  396 + # Setup times container
  397 + times = self._setup_times()
  398 +
  399 + # Test times[start:end]
  400 + self.test_value(len(times[3:5]), 2)
  401 + self.test_value(times[3:5][0].secs(), 3.0)
  402 + self.test_value(times[3:5][1].secs(), 4.0)
  403 +
  404 + # Test times[start:]
  405 + self.test_value(len(times[7:]), 3)
  406 + self.test_value(times[7:][0].secs(), 7.0)
  407 + self.test_value(times[7:][1].secs(), 8.0)
  408 + self.test_value(times[7:][2].secs(), 9.0)
  409 +
  410 + # Test times[:end]
  411 + self.test_value(len(times[:2]), 2)
  412 + self.test_value(times[:2][0].secs(), 0.0)
  413 + self.test_value(times[:2][1].secs(), 1.0)
  414 +
  415 + # Test times[:]
  416 + self.test_value(len(times[:]), 10)
  417 + for i in range(10):
  418 + self.test_value(times[:][i].secs(), float(i))
  419 +
  420 + # Test times[start:end:step]
  421 + self.test_value(len(times[3:7:2]), 2)
  422 + self.test_value(times[3:7:2][0].secs(), 3.0)
  423 + self.test_value(times[3:7:2][1].secs(), 5.0)
  424 +
  425 + # Test times[start:end:step]
  426 + self.test_value(len(times[6:3:-2]), 2)
  427 + self.test_value(times[6:3:-2][0].secs(), 6.0)
  428 + self.test_value(times[6:3:-2][1].secs(), 4.0)
  429 +
  430 + # Test times[-start:]
  431 + self.test_value(len(times[-2:]), 2)
  432 + self.test_value(times[-2:][0].secs(), 8.0)
  433 + self.test_value(times[-2:][1].secs(), 9.0)
  434 +
  435 + # Test times[:-end]
  436 + self.test_value(len(times[:-7]), 3)
  437 + self.test_value(times[:-7][0].secs(), 0.0)
  438 + self.test_value(times[:-7][1].secs(), 1.0)
  439 + self.test_value(times[:-7][2].secs(), 2.0)
  440 +
  441 + # Return
  442 + return
  443 +
  444 + # Test GPhotons class access operators
  445 + def _test_photons_access(self):
  446 + """
  447 + Test GPhotons class observation access
  448 + """
  449 + # Setup photons container
  450 + photons = self._setup_photons()
  451 +
  452 + # Loop over all elements using the container iterator and count the
  453 + # number of iterations
  454 + nphotons = 0
  455 + reference = 0.0
  456 + for photon in photons:
  457 + self.test_value(photon.energy().MeV(), reference)
  458 + nphotons += 1
  459 + reference += 1.0
  460 +
  461 + # Check that looping was successful
  462 + self.test_value(photons.size(), 10, 'Check container size')
  463 + self.test_value(nphotons, 10, 'Check container iterator')
  464 +
  465 + # Test access from start
  466 + self.test_value(photons[3].energy().MeV(), 3.0)
  467 +
  468 + # Check access from end
  469 + self.test_value(photons[-2].energy().MeV(), 8.0)
  470 +
  471 + # Return
  472 + return
  473 +
  474 + # Test GPhotons class slicing
  475 + def _test_photons_slicing(self):
  476 + """
  477 + Test GPhotons class slicing
  478 + """
  479 + # Setup photons container
  480 + photons = self._setup_photons()
  481 +
  482 + # Test photons[start:end]
  483 + self.test_value(len(photons[3:5]), 2)
  484 + self.test_value(photons[3:5][0].energy().MeV(), 3.0)
  485 + self.test_value(photons[3:5][1].energy().MeV(), 4.0)
  486 +
  487 + # Test photons[start:]
  488 + self.test_value(len(photons[7:]), 3)
  489 + self.test_value(photons[7:][0].energy().MeV(), 7.0)
  490 + self.test_value(photons[7:][1].energy().MeV(), 8.0)
  491 + self.test_value(photons[7:][2].energy().MeV(), 9.0)
  492 +
  493 + # Test photons[:end]
  494 + self.test_value(len(photons[:2]), 2)
  495 + self.test_value(photons[:2][0].energy().MeV(), 0.0)
  496 + self.test_value(photons[:2][1].energy().MeV(), 1.0)
  497 +
  498 + # Test photons[:]
  499 + self.test_value(len(photons[:]), 10)
  500 + for i in range(10):
  501 + self.test_value(photons[:][i].energy().MeV(), float(i))
  502 +
  503 + # Test photons[start:end:step]
  504 + self.test_value(len(photons[3:7:2]), 2)
  505 + self.test_value(photons[3:7:2][0].energy().MeV(), 3.0)
  506 + self.test_value(photons[3:7:2][1].energy().MeV(), 5.0)
  507 +
  508 + # Test photons[start:end:step]
  509 + self.test_value(len(photons[6:3:-2]), 2)
  510 + self.test_value(photons[6:3:-2][0].energy().MeV(), 6.0)
  511 + self.test_value(photons[6:3:-2][1].energy().MeV(), 4.0)
  512 +
  513 + # Test photons[-start:]
  514 + self.test_value(len(photons[-2:]), 2)
  515 + self.test_value(photons[-2:][0].energy().MeV(), 8.0)
  516 + self.test_value(photons[-2:][1].energy().MeV(), 9.0)
  517 +
  518 + # Test photons[:-end]
  519 + self.test_value(len(photons[:-7]), 3)
  520 + self.test_value(photons[:-7][0].energy().MeV(), 0.0)
  521 + self.test_value(photons[:-7][1].energy().MeV(), 1.0)
  522 + self.test_value(photons[:-7][2].energy().MeV(), 2.0)
  523 +
  524 + # Return
  525 + return
  526 +
234 527 # Test GObservations class access operators
235 528 def _test_observations_access(self):
236 529 """
... ... @@ -289,11 +582,17 @@ class Test(gammalib.GPythonTestSuite):
289 582 self.name('obs')
290 583  
291 584 # Append tests
292   - self.append(self._test_energy, 'Test GEnergy class')
293   - self.append(self._test_energies, 'Test GEnergies class')
294   - self.append(self._test_time, 'Test GTime class')
295   - self.append(self._test_observations_access, 'Test GObservations class observation access')
296   - self.append(self._test_observations_slicing, 'Test GObservations class slicing')
  585 + self.append(self._test_energy, 'Test GEnergy')
  586 + self.append(self._test_energies, 'Test GEnergies')
  587 + self.append(self._test_energies_access, 'Test GEnergies energy access')
  588 + self.append(self._test_energies_slicing, 'Test GEnergies slicing')
  589 + self.append(self._test_time, 'Test GTime')
  590 + self.append(self._test_times_access, 'Test GTimes time access')
  591 + self.append(self._test_times_slicing, 'Test GTimes slicing')
  592 + self.append(self._test_photons_access, 'Test GPhotons photon access')
  593 + self.append(self._test_photons_slicing, 'Test GPhotons slicing')
  594 + self.append(self._test_observations_access, 'Test GObservations observation access')
  595 + self.append(self._test_observations_slicing, 'Test GObservations slicing')
297 596  
298 597 # Return
299 598 return
... ...