Simplified the usage of write_atomic_weights.py