0

I'm very new to Python, so please forgive my ignorance. I'm trying to calculate the total number of energy units in a system. For example, the Omega here will output both (0,0,0,1) and (2,2,2,1) along with a whole lot of other tuples. I want to extract from Omega how many tuples have a total value of 1 (like the first example) and how many have a total value of 7 (like the second example). How do I achieve this?

import numpy as np
import matplotlib.pyplot as plt
from itertools import product

N = 4 ##The number of Oscillators
q = range(3) ## Range of number of possible energy units per oscillator


Omega = product(q, repeat = N)
print(list(product(q, repeat = N)))
MarianD
  • 13,096
  • 12
  • 42
  • 54
Jonathan Schoenmaker
  • 1
  • 2
    you can use `filter` or list comprehension to achieve this, for example, `[tup for tup in Omega if sum(tup) == 1]` – gold_cy May 03 '21 at 11:19

2 Answers2

1

try this:

Omega = product(q, repeat = N)
l = list(product(q, repeat = N))
l1 = [i for i in l if sum(i)==1]
l2 = [i for i in l if sum(i)==7]
print(l1,l2)
Amit Nanaware
  • 3,203
  • 1
  • 6
  • 19
  • You never used the first line. The 2nd line is superfluous - you may use directly `Omega` instead of `l` in the next 2 lines. BTW, the name `Omega` (with capital `O`) doesn't conform [PEP 8 - Style Guide for Python Code](https://www.python.org/dev/peps/pep-0008/#naming-conventions). Use `omega` instead. – MarianD May 03 '21 at 11:53
0

I believe you can use sum() on tuples as well as lists of integers/numbers.

Now you say omega is a list of tuples, is that correct? Something like

Omega = [(0,0,0,1), (2,2,2,1), ...)]

In that case I think you can do

sums_to_1 = [int_tuple for int_tuple in omega if sum(int_tuple) == 1]

If you want to have some default value for the tuples that don't sum to one you can put the if statement in the list comprehension in the beginning and do

sums_to_1 = [int_tuple if sum(int_tuple) == 1 else 'SomeDefaultValue' for int_tuple in omega]

Brian Karabinchak
  • 207
  • 1
  • 9
  • I checked it out and it's true. Apparnetly importing numpy certain ways overwrites the sum() function though I don't believe op needs to worry about that unless he does have a `from numpy import *` somewhere in which case sum can get funky. https://stackoverflow.com/questions/39552458/python-sum-has-a-different-result-after-importing-numpy – Brian Karabinchak May 03 '21 at 11:22
  • 1
    Then you will give the OP greater confidence in your answer if you omit the "believe" and "think". That is what I think, anyway. – Booboo May 03 '21 at 11:26
  • You probably don't try to run your code. If you fix it (change the first `Omega` to `omega`) and remove the phrases "I believe" and "I think", I'll vote up your answer. – MarianD May 03 '21 at 12:02